chore: refactor the imports (#2406)

* feat: mysql and redis metric support

* feat: mysql and redis metric support

* feat: mysql and redis metric support

Co-authored-by: dawn.zhou <dawn.zhou@yijinin.com>

.github/dependabot.yml

@@ -0,0 +1,11 @@
+# To get started with Dependabot version updates, you'll need to specify which
+# package ecosystems to update and where the package manifests are located.
+# Please see the documentation for all configuration options:
+# https://docs.github.com/github/administering-a-repository/configuration-options-for-dependency-updates
+
+version: 2
+updates:
+  - package-ecosystem: "gomod" # See documentation for possible values
+    directory: "/" # Location of package manifests
+    schedule:
+      interval: "daily"

.github/workflows/codeql-analysis.yml

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

.github/workflows/go.yml

@@ -14,11 +14,11 @@ jobs:
       - name: Set up Go 1.x
         uses: actions/setup-go@v2
         with:
-          go-version: ^1.15
+          go-version: ^1.16
         id: go
 
       - name: Check out code into the Go module directory
-        uses: actions/checkout@v2
+        uses: actions/checkout@v3
 
       - name: Get dependencies
         run: |
@@ -43,10 +43,10 @@ jobs:
       - name: Set up Go 1.x
         uses: actions/setup-go@v2
         with:
-          go-version: ^1.15
+          go-version: ^1.16
 
       - name: Checkout codebase
-        uses: actions/checkout@v2
+        uses: actions/checkout@v3
 
       - name: Test
         run: |

.github/workflows/issues.yml

@@ -9,8 +9,8 @@ jobs:
     steps:
       - uses: actions/stale@v3
         with:
-          days-before-issue-stale: 30
-          days-before-issue-close: 14
+          days-before-issue-stale: 365
+          days-before-issue-close: 90
           stale-issue-label: "stale"
           stale-issue-message: "This issue is stale because it has been open for 30 days with no activity."
           close-issue-message: "This issue was closed because it has been inactive for 14 days since being marked as stale."

.github/workflows/release.yaml

@@ -16,7 +16,7 @@ jobs:
           - goarch: "386"
             goos: darwin
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
       - uses: zeromicro/go-zero-release-action@master
         with:
           github_token: ${{ secrets.GITHUB_TOKEN }}
@@ -25,4 +25,4 @@ jobs:
           goversion: "https://dl.google.com/go/go1.17.5.linux-amd64.tar.gz"
           project_path: "tools/goctl"
           binary_name: "goctl"
-          extra_files: tools/goctl/goctl.md
+          extra_files: tools/goctl/readme.md tools/goctl/readme-cn.md

.github/workflows/reviewdog.yml

@@ -5,7 +5,7 @@ jobs:
     name: runner / staticcheck
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
       - uses: reviewdog/action-staticcheck@v1
         with:
           github_token: ${{ secrets.github_token }}

core/codec/rsa.go

@@ -7,7 +7,7 @@ import (
 	"encoding/base64"
 	"encoding/pem"
 	"errors"
-	"io/ioutil"
+	"os"
 )
 
 var (
@@ -48,7 +48,7 @@ type (
 
 // NewRsaDecrypter returns a RsaDecrypter with the given file.
 func NewRsaDecrypter(file string) (RsaDecrypter, error) {
-	content, err := ioutil.ReadFile(file)
+	content, err := os.ReadFile(file)
 	if err != nil {
 		return nil, err
 	}

core/collection/cache.go

@@ -26,7 +26,7 @@ type (
 	// CacheOption defines the method to customize a Cache.
 	CacheOption func(cache *Cache)
 
-	// A Cache object is a in-memory cache.
+	// A Cache object is an in-memory cache.
 	Cache struct {
 		name           string
 		lock           sync.Mutex

core/collection/safemap.go

@@ -68,6 +68,24 @@ func (m *SafeMap) Get(key interface{}) (interface{}, bool) {
 	return val, ok
 }
 
+// Range calls f sequentially for each key and value present in the map.
+// If f returns false, range stops the iteration.
+func (m *SafeMap) Range(f func(key, val interface{}) bool) {
+	m.lock.RLock()
+	defer m.lock.RUnlock()
+
+	for k, v := range m.dirtyOld {
+		if !f(k, v) {
+			return
+		}
+	}
+	for k, v := range m.dirtyNew {
+		if !f(k, v) {
+			return
+		}
+	}
+}
+
 // Set sets the value into m with the given key.
 func (m *SafeMap) Set(key, value interface{}) {
 	m.lock.Lock()

core/collection/safemap_test.go

@@ -1,6 +1,7 @@
 package collection
 
 import (
+	"sync/atomic"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -107,3 +108,42 @@ func testSafeMapWithParameters(t *testing.T, size, exception int) {
 		}
 	}
 }
+
+func TestSafeMap_Range(t *testing.T) {
+	const (
+		size       = 100000
+		exception1 = 5
+		exception2 = 500
+	)
+
+	m := NewSafeMap()
+	newMap := NewSafeMap()
+
+	for i := 0; i < size; i++ {
+		m.Set(i, i)
+	}
+	for i := 0; i < size; i++ {
+		if i%exception1 == 0 {
+			m.Del(i)
+		}
+	}
+
+	for i := size; i < size<<1; i++ {
+		m.Set(i, i)
+	}
+	for i := size; i < size<<1; i++ {
+		if i%exception2 != 0 {
+			m.Del(i)
+		}
+	}
+
+	var count int32
+	m.Range(func(k, v interface{}) bool {
+		atomic.AddInt32(&count, 1)
+		newMap.Set(k, v)
+		return true
+	})
+	assert.Equal(t, int(atomic.LoadInt32(&count)), m.Size())
+	assert.Equal(t, m.dirtyNew, newMap.dirtyNew)
+	assert.Equal(t, m.dirtyOld, newMap.dirtyOld)
+}

core/conf/config.go

@@ -2,7 +2,6 @@ package conf
 
 import (
 	"fmt"
-	"io/ioutil"
 	"log"
 	"os"
 	"path"
@@ -20,7 +19,7 @@ var loaders = map[string]func([]byte, interface{}) error{
 
 // Load loads config into v from file, .json, .yaml and .yml are acceptable.
 func Load(file string, v interface{}, opts ...Option) error {
-	content, err := ioutil.ReadFile(file)
+	content, err := os.ReadFile(file)
 	if err != nil {
 		return err
 	}

core/conf/config_test.go

@@ -1,7 +1,6 @@
 package conf
 
 import (
-	"io/ioutil"
 	"os"
 	"testing"
 
@@ -146,12 +145,12 @@ func TestConfigJsonEnv(t *testing.T) {
 }
 
 func createTempFile(ext, text string) (string, error) {
-	tmpfile, err := ioutil.TempFile(os.TempDir(), hash.Md5Hex([]byte(text))+"*"+ext)
+	tmpfile, err := os.CreateTemp(os.TempDir(), hash.Md5Hex([]byte(text))+"*"+ext)
 	if err != nil {
 		return "", err
 	}
 
-	if err := ioutil.WriteFile(tmpfile.Name(), []byte(text), os.ModeTemporary); err != nil {
+	if err := os.WriteFile(tmpfile.Name(), []byte(text), os.ModeTemporary); err != nil {
 		return "", err
 	}
 

core/conf/readme.md

@@ -3,10 +3,10 @@
 1. Define a config structure, like below:
 
 ```go
-RestfulConf struct {
+type RestfulConf struct {
 	Host         string        `json:",default=0.0.0.0"`
 	Port         int
-	LogMode      string        `json:",options=[file,console]"
+	LogMode      string        `json:",options=[file,console]"`
 	Verbose      bool          `json:",optional"`
 	MaxConns     int           `json:",default=10000"`
 	MaxBytes     int64         `json:",default=1048576"`
@@ -15,7 +15,9 @@ RestfulConf struct {
 }
 ```
 
-2. Write the yaml or json config file:
+2. Write the yaml, toml or json config file:
+
+- yaml example
 
 ```yaml
 # most fields are optional or have default values
@@ -25,6 +27,16 @@ LogMode: console
 MaxBytes: ${MAX_BYTES}
 ```
 
+- toml example
+
+```toml
+# most fields are optional or have default values
+Port = 8_080
+LogMode = "console"
+# you can use env settings
+MaxBytes = "${MAX_BYTES}"
+```
+
 3. Load the config from a file:
 
 ```go

core/discov/internal/accountmanager.go

@@ -3,7 +3,7 @@ package internal
 import (
 	"crypto/tls"
 	"crypto/x509"
-	"io/ioutil"
+	"os"
 	"sync"
 )
 
@@ -37,7 +37,7 @@ func AddTLS(endpoints []string, certFile, certKeyFile, caFile string, insecureSk
 		return err
 	}
 
-	caData, err := ioutil.ReadFile(caFile)
+	caData, err := os.ReadFile(caFile)
 	if err != nil {
 		return err
 	}

core/discov/internal/registry.go

@@ -191,9 +191,11 @@ func (c *cluster) handleWatchEvents(key string, events []*clientv3.Event) {
 				})
 			}
 		case clientv3.EventTypeDelete:
+			c.lock.Lock()
 			if vals, ok := c.values[key]; ok {
 				delete(vals, string(ev.Kv.Key))
 			}
+			c.lock.Unlock()
 			for _, l := range listeners {
 				l.OnDelete(KV{
 					Key: string(ev.Kv.Key),
@@ -206,7 +208,7 @@ func (c *cluster) handleWatchEvents(key string, events []*clientv3.Event) {
 	}
 }
 
-func (c *cluster) load(cli EtcdClient, key string) {
+func (c *cluster) load(cli EtcdClient, key string) int64 {
 	var resp *clientv3.GetResponse
 	for {
 		var err error
@@ -230,6 +232,8 @@ func (c *cluster) load(cli EtcdClient, key string) {
 	}
 
 	c.handleChanges(key, kvs)
+
+	return resp.Header.Revision
 }
 
 func (c *cluster) monitor(key string, l UpdateListener) error {
@@ -242,9 +246,9 @@ func (c *cluster) monitor(key string, l UpdateListener) error {
 		return err
 	}
 
-	c.load(cli, key)
+	rev := c.load(cli, key)
 	c.watchGroup.Run(func() {
-		c.watch(cli, key)
+		c.watch(cli, key, rev)
 	})
 
 	return nil
@@ -276,22 +280,28 @@ func (c *cluster) reload(cli EtcdClient) {
 	for _, key := range keys {
 		k := key
 		c.watchGroup.Run(func() {
-			c.load(cli, k)
-			c.watch(cli, k)
+			rev := c.load(cli, k)
+			c.watch(cli, k, rev)
 		})
 	}
 }
 
-func (c *cluster) watch(cli EtcdClient, key string) {
+func (c *cluster) watch(cli EtcdClient, key string, rev int64) {
 	for {
-		if c.watchStream(cli, key) {
+		if c.watchStream(cli, key, rev) {
 			return
 		}
 	}
 }
 
-func (c *cluster) watchStream(cli EtcdClient, key string) bool {
-	rch := cli.Watch(clientv3.WithRequireLeader(c.context(cli)), makeKeyPrefix(key), clientv3.WithPrefix())
+func (c *cluster) watchStream(cli EtcdClient, key string, rev int64) bool {
+	var rch clientv3.WatchChan
+	if rev != 0 {
+		rch = cli.Watch(clientv3.WithRequireLeader(c.context(cli)), makeKeyPrefix(key), clientv3.WithPrefix(), clientv3.WithRev(rev+1))
+	} else {
+		rch = cli.Watch(clientv3.WithRequireLeader(c.context(cli)), makeKeyPrefix(key), clientv3.WithPrefix())
+	}
+
 	for {
 		select {
 		case wresp, ok := <-rch:

core/discov/internal/registry_test.go

@@ -2,6 +2,7 @@ package internal
 
 import (
 	"context"
+	"go.etcd.io/etcd/api/v3/etcdserverpb"
 	"sync"
 	"testing"
 
@@ -112,6 +113,7 @@ func TestCluster_Load(t *testing.T) {
 	restore := setMockClient(cli)
 	defer restore()
 	cli.EXPECT().Get(gomock.Any(), "any/", gomock.Any()).Return(&clientv3.GetResponse{
+		Header: &etcdserverpb.ResponseHeader{},
 		Kvs: []*mvccpb.KeyValue{
 			{
 				Key:   []byte("hello"),
@@ -168,7 +170,7 @@ func TestCluster_Watch(t *testing.T) {
 			listener.EXPECT().OnDelete(gomock.Any()).Do(func(_ interface{}) {
 				wg.Done()
 			}).MaxTimes(1)
-			go c.watch(cli, "any")
+			go c.watch(cli, "any", 0)
 			ch <- clientv3.WatchResponse{
 				Events: []*clientv3.Event{
 					{
@@ -212,7 +214,7 @@ func TestClusterWatch_RespFailures(t *testing.T) {
 				ch <- resp
 				close(c.done)
 			}()
-			c.watch(cli, "any")
+			c.watch(cli, "any", 0)
 		})
 	}
 }
@@ -232,7 +234,7 @@ func TestClusterWatch_CloseChan(t *testing.T) {
 		close(ch)
 		close(c.done)
 	}()
-	c.watch(cli, "any")
+	c.watch(cli, "any", 0)
 }
 
 func TestValueOnlyContext(t *testing.T) {

core/errorx/wrap.go

@@ -0,0 +1,21 @@
+package errorx
+
+import "fmt"
+
+// Wrap returns an error that wraps err with given message.
+func Wrap(err error, message string) error {
+	if err == nil {
+		return nil
+	}
+
+	return fmt.Errorf("%s: %w", message, err)
+}
+
+// Wrapf returns an error that wraps err with given format and args.
+func Wrapf(err error, format string, args ...interface{}) error {
+	if err == nil {
+		return nil
+	}
+
+	return fmt.Errorf("%s: %w", fmt.Sprintf(format, args...), err)
+}

core/errorx/wrap_test.go

@@ -0,0 +1,24 @@
+package errorx
+
+import (
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestWrap(t *testing.T) {
+	assert.Nil(t, Wrap(nil, "test"))
+	assert.Equal(t, "foo: bar", Wrap(errors.New("bar"), "foo").Error())
+
+	err := errors.New("foo")
+	assert.True(t, errors.Is(Wrap(err, "bar"), err))
+}
+
+func TestWrapf(t *testing.T) {
+	assert.Nil(t, Wrapf(nil, "%s", "test"))
+	assert.Equal(t, "foo bar: quz", Wrapf(errors.New("quz"), "foo %s", "bar").Error())
+
+	err := errors.New("foo")
+	assert.True(t, errors.Is(Wrapf(err, "foo %s", "bar"), err))
+}

core/fs/temps.go

@@ -1,7 +1,6 @@
 package fs
 
 import (
-	"io/ioutil"
 	"os"
 
 	"github.com/zeromicro/go-zero/core/hash"
@@ -12,12 +11,12 @@ import (
 // The file is kept as open, the caller should close the file handle,
 // and remove the file by name.
 func TempFileWithText(text string) (*os.File, error) {
-	tmpfile, err := ioutil.TempFile(os.TempDir(), hash.Md5Hex([]byte(text)))
+	tmpfile, err := os.CreateTemp(os.TempDir(), hash.Md5Hex([]byte(text)))
 	if err != nil {
 		return nil, err
 	}
 
-	if err := ioutil.WriteFile(tmpfile.Name(), []byte(text), os.ModeTemporary); err != nil {
+	if err := os.WriteFile(tmpfile.Name(), []byte(text), os.ModeTemporary); err != nil {
 		return nil, err
 	}
 

core/fs/temps_test.go

@@ -1,7 +1,7 @@
 package fs
 
 import (
-	"io/ioutil"
+	"io"
 	"os"
 	"testing"
 
@@ -21,7 +21,7 @@ func TestTempFileWithText(t *testing.T) {
 	}
 	defer os.Remove(f.Name())
 
-	bs, err := ioutil.ReadAll(f)
+	bs, err := io.ReadAll(f)
 	assert.Nil(t, err)
 	if len(bs) != 4 {
 		t.Error("TempFileWithText returned wrong file size")
@@ -41,7 +41,7 @@ func TestTempFilenameWithText(t *testing.T) {
 	}
 	defer os.Remove(f)
 
-	bs, err := ioutil.ReadFile(f)
+	bs, err := os.ReadFile(f)
 	assert.Nil(t, err)
 	if len(bs) != 4 {
 		t.Error("TempFilenameWithText returned wrong file size")

core/fx/stream_test.go

@@ -1,7 +1,7 @@
 package fx
 
 import (
-	"io/ioutil"
+	"io"
 	"log"
 	"math/rand"
 	"reflect"
@@ -238,7 +238,7 @@ func TestLast(t *testing.T) {
 
 func TestMap(t *testing.T) {
 	runCheckedTest(t, func(t *testing.T) {
-		log.SetOutput(ioutil.Discard)
+		log.SetOutput(io.Discard)
 
 		tests := []struct {
 			mapper MapFunc

core/iox/read.go

@@ -4,7 +4,6 @@ import (
 	"bufio"
 	"bytes"
 	"io"
-	"io/ioutil"
 	"os"
 	"strings"
 )
@@ -26,7 +25,7 @@ type (
 func DupReadCloser(reader io.ReadCloser) (io.ReadCloser, io.ReadCloser) {
 	var buf bytes.Buffer
 	tee := io.TeeReader(reader, &buf)
-	return ioutil.NopCloser(tee), ioutil.NopCloser(&buf)
+	return io.NopCloser(tee), io.NopCloser(&buf)
 }
 
 // KeepSpace customizes the reading functions to keep leading and tailing spaces.
@@ -54,7 +53,7 @@ func ReadBytes(reader io.Reader, buf []byte) error {
 
 // ReadText reads content from the given file with leading and tailing spaces trimmed.
 func ReadText(filename string) (string, error) {
-	content, err := ioutil.ReadFile(filename)
+	content, err := os.ReadFile(filename)
 	if err != nil {
 		return "", err
 	}

core/iox/read_test.go

@@ -3,7 +3,6 @@ package iox
 import (
 	"bytes"
 	"io"
-	"io/ioutil"
 	"os"
 	"testing"
 	"time"
@@ -97,10 +96,10 @@ func TestReadTextLines(t *testing.T) {
 
 func TestDupReadCloser(t *testing.T) {
 	input := "hello"
-	reader := ioutil.NopCloser(bytes.NewBufferString(input))
+	reader := io.NopCloser(bytes.NewBufferString(input))
 	r1, r2 := DupReadCloser(reader)
 	verify := func(r io.Reader) {
-		output, err := ioutil.ReadAll(r)
+		output, err := io.ReadAll(r)
 		assert.Nil(t, err)
 		assert.Equal(t, input, string(output))
 	}
@@ -110,7 +109,7 @@ func TestDupReadCloser(t *testing.T) {
 }
 
 func TestReadBytes(t *testing.T) {
-	reader := ioutil.NopCloser(bytes.NewBufferString("helloworld"))
+	reader := io.NopCloser(bytes.NewBufferString("helloworld"))
 	buf := make([]byte, 5)
 	err := ReadBytes(reader, buf)
 	assert.Nil(t, err)
@@ -118,7 +117,7 @@ func TestReadBytes(t *testing.T) {
 }
 
 func TestReadBytesNotEnough(t *testing.T) {
-	reader := ioutil.NopCloser(bytes.NewBufferString("hell"))
+	reader := io.NopCloser(bytes.NewBufferString("hell"))
 	buf := make([]byte, 5)
 	err := ReadBytes(reader, buf)
 	assert.Equal(t, io.EOF, err)

core/iox/textfile_test.go

@@ -1,7 +1,6 @@
 package iox
 
 import (
-	"io/ioutil"
 	"os"
 	"testing"
 
@@ -13,7 +12,7 @@ func TestCountLines(t *testing.T) {
 2
 3
 4`
-	file, err := ioutil.TempFile(os.TempDir(), "test-")
+	file, err := os.CreateTemp(os.TempDir(), "test-")
 	if err != nil {
 		t.Fatal(err)
 	}

core/jsonx/json.go

@@ -13,6 +13,16 @@ func Marshal(v interface{}) ([]byte, error) {
 	return json.Marshal(v)
 }
 
+// MarshalToString marshals v into a string.
+func MarshalToString(v interface{}) (string, error) {
+	data, err := Marshal(v)
+	if err != nil {
+		return "", err
+	}
+
+	return string(data), nil
+}
+
 // Unmarshal unmarshals data bytes into v.
 func Unmarshal(data []byte, v interface{}) error {
 	decoder := json.NewDecoder(bytes.NewReader(data))

core/jsonx/json_test.go

@@ -20,6 +20,22 @@ func TestMarshal(t *testing.T) {
 	assert.Equal(t, `{"name":"John","age":30}`, string(bs))
 }
 
+func TestMarshalToString(t *testing.T) {
+	var v = struct {
+		Name string `json:"name"`
+		Age  int    `json:"age"`
+	}{
+		Name: "John",
+		Age:  30,
+	}
+	toString, err := MarshalToString(v)
+	assert.Nil(t, err)
+	assert.Equal(t, `{"name":"John","age":30}`, toString)
+
+	_, err = MarshalToString(make(chan int))
+	assert.NotNil(t, err)
+}
+
 func TestUnmarshal(t *testing.T) {
 	const s = `{"name":"John","age":30}`
 	var v struct {

core/limit/periodlimit.go

@@ -1,6 +1,7 @@
 package limit
 
 import (
+	"context"
 	"errors"
 	"strconv"
 	"time"
@@ -74,7 +75,12 @@ func NewPeriodLimit(period, quota int, limitStore *redis.Redis, keyPrefix string
 
 // Take requests a permit, it returns the permit state.
 func (h *PeriodLimit) Take(key string) (int, error) {
-	resp, err := h.limitStore.Eval(periodScript, []string{h.keyPrefix + key}, []string{
+	return h.TakeCtx(context.Background(), key)
+}
+
+// TakeCtx requests a permit with context, it returns the permit state.
+func (h *PeriodLimit) TakeCtx(ctx context.Context, key string) (int, error) {
+	resp, err := h.limitStore.EvalCtx(ctx, periodScript, []string{h.keyPrefix + key}, []string{
 		strconv.Itoa(h.quota),
 		strconv.Itoa(h.calcExpireSeconds()),
 	})

core/logx/config.go

@@ -11,4 +11,16 @@ type LogConf struct {
 	Compress            bool   `json:",optional"`
 	KeepDays            int    `json:",optional"`
 	StackCooldownMillis int    `json:",default=100"`
+	// MaxBackups represents how many backup log files will be kept. 0 means all files will be kept forever.
+	// Only take effect when RotationRuleType is `size`.
+	// Even thougth `MaxBackups` sets 0, log files will still be removed
+	// if the `KeepDays` limitation is reached.
+	MaxBackups int `json:",default=0"`
+	// MaxSize represents how much space the writing log file takes up. 0 means no limit. The unit is `MB`.
+	// Only take effect when RotationRuleType is `size`
+	MaxSize int `json:",default=0"`
+	// RotationRuleType represents the type of log rotation rule. Default is `daily`.
+	// daily: daily rotation.
+	// size: size limited rotation.
+	Rotation string `json:",default=daily,options=[daily,size]"`
 }

core/logx/contextlogger.go

@@ -11,65 +11,65 @@ import (
 
 // WithContext sets ctx to log, for keeping tracing information.
 func WithContext(ctx context.Context) Logger {
-	return &traceLogger{
+	return &contextLogger{
 		ctx: ctx,
 	}
 }
 
-type traceLogger struct {
+type contextLogger struct {
 	logEntry
 	ctx context.Context
 }
 
-func (l *traceLogger) Error(v ...interface{}) {
+func (l *contextLogger) Error(v ...interface{}) {
 	l.err(fmt.Sprint(v...))
 }
 
-func (l *traceLogger) Errorf(format string, v ...interface{}) {
+func (l *contextLogger) Errorf(format string, v ...interface{}) {
 	l.err(fmt.Sprintf(format, v...))
 }
 
-func (l *traceLogger) Errorv(v interface{}) {
+func (l *contextLogger) Errorv(v interface{}) {
 	l.err(fmt.Sprint(v))
 }
 
-func (l *traceLogger) Errorw(msg string, fields ...LogField) {
+func (l *contextLogger) Errorw(msg string, fields ...LogField) {
 	l.err(msg, fields...)
 }
 
-func (l *traceLogger) Info(v ...interface{}) {
+func (l *contextLogger) Info(v ...interface{}) {
 	l.info(fmt.Sprint(v...))
 }
 
-func (l *traceLogger) Infof(format string, v ...interface{}) {
+func (l *contextLogger) Infof(format string, v ...interface{}) {
 	l.info(fmt.Sprintf(format, v...))
 }
 
-func (l *traceLogger) Infov(v interface{}) {
+func (l *contextLogger) Infov(v interface{}) {
 	l.info(v)
 }
 
-func (l *traceLogger) Infow(msg string, fields ...LogField) {
+func (l *contextLogger) Infow(msg string, fields ...LogField) {
 	l.info(msg, fields...)
 }
 
-func (l *traceLogger) Slow(v ...interface{}) {
+func (l *contextLogger) Slow(v ...interface{}) {
 	l.slow(fmt.Sprint(v...))
 }
 
-func (l *traceLogger) Slowf(format string, v ...interface{}) {
+func (l *contextLogger) Slowf(format string, v ...interface{}) {
 	l.slow(fmt.Sprintf(format, v...))
 }
 
-func (l *traceLogger) Slowv(v interface{}) {
+func (l *contextLogger) Slowv(v interface{}) {
 	l.slow(v)
 }
 
-func (l *traceLogger) Sloww(msg string, fields ...LogField) {
+func (l *contextLogger) Sloww(msg string, fields ...LogField) {
 	l.slow(msg, fields...)
 }
 
-func (l *traceLogger) WithContext(ctx context.Context) Logger {
+func (l *contextLogger) WithContext(ctx context.Context) Logger {
 	if ctx == nil {
 		return l
 	}
@@ -78,40 +78,49 @@ func (l *traceLogger) WithContext(ctx context.Context) Logger {
 	return l
 }
 
-func (l *traceLogger) WithDuration(duration time.Duration) Logger {
+func (l *contextLogger) WithDuration(duration time.Duration) Logger {
 	l.Duration = timex.ReprOfDuration(duration)
 	return l
 }
 
-func (l *traceLogger) buildFields(fields ...LogField) []LogField {
+func (l *contextLogger) buildFields(fields ...LogField) []LogField {
 	if len(l.Duration) > 0 {
 		fields = append(fields, Field(durationKey, l.Duration))
 	}
+
 	traceID := traceIdFromContext(l.ctx)
 	if len(traceID) > 0 {
 		fields = append(fields, Field(traceKey, traceID))
 	}
+
 	spanID := spanIdFromContext(l.ctx)
 	if len(spanID) > 0 {
 		fields = append(fields, Field(spanKey, spanID))
 	}
 
+	val := l.ctx.Value(fieldsContextKey)
+	if val != nil {
+		if arr, ok := val.([]LogField); ok {
+			fields = append(fields, arr...)
+		}
+	}
+
 	return fields
 }
 
-func (l *traceLogger) err(v interface{}, fields ...LogField) {
+func (l *contextLogger) err(v interface{}, fields ...LogField) {
 	if shallLog(ErrorLevel) {
 		getWriter().Error(v, l.buildFields(fields...)...)
 	}
 }
 
-func (l *traceLogger) info(v interface{}, fields ...LogField) {
+func (l *contextLogger) info(v interface{}, fields ...LogField) {
 	if shallLog(InfoLevel) {
 		getWriter().Info(v, l.buildFields(fields...)...)
 	}
 }
 
-func (l *traceLogger) slow(v interface{}, fields ...LogField) {
+func (l *contextLogger) slow(v interface{}, fields ...LogField) {
 	if shallLog(ErrorLevel) {
 		getWriter().Slow(v, l.buildFields(fields...)...)
 	}

core/logx/contextlogger_test.go

@@ -29,7 +29,7 @@ func TestTraceLog(t *testing.T) {
 	otel.SetTracerProvider(tp)
 	defer otel.SetTracerProvider(otp)
 
-	ctx, span := tp.Tracer("foo").Start(context.Background(), "bar")
+	ctx, span := tp.Tracer("trace-id").Start(context.Background(), "span-id")
 	defer span.End()
 
 	WithContext(ctx).Info(testlog)
@@ -50,7 +50,7 @@ func TestTraceError(t *testing.T) {
 	otel.SetTracerProvider(tp)
 	defer otel.SetTracerProvider(otp)
 
-	ctx, span := tp.Tracer("foo").Start(context.Background(), "bar")
+	ctx, span := tp.Tracer("trace-id").Start(context.Background(), "span-id")
 	defer span.End()
 
 	var nilCtx context.Context
@@ -67,10 +67,10 @@ func TestTraceError(t *testing.T) {
 	l.WithDuration(time.Second).Errorv(testlog)
 	validate(t, w.String(), true, true)
 	w.Reset()
-	l.WithDuration(time.Second).Errorw(testlog, Field("foo", "bar"))
+	l.WithDuration(time.Second).Errorw(testlog, Field("basket", "ball"))
 	validate(t, w.String(), true, true)
-	assert.True(t, strings.Contains(w.String(), "foo"), w.String())
-	assert.True(t, strings.Contains(w.String(), "bar"), w.String())
+	assert.True(t, strings.Contains(w.String(), "basket"), w.String())
+	assert.True(t, strings.Contains(w.String(), "ball"), w.String())
 }
 
 func TestTraceInfo(t *testing.T) {
@@ -87,7 +87,7 @@ func TestTraceInfo(t *testing.T) {
 	otel.SetTracerProvider(tp)
 	defer otel.SetTracerProvider(otp)
 
-	ctx, span := tp.Tracer("foo").Start(context.Background(), "bar")
+	ctx, span := tp.Tracer("trace-id").Start(context.Background(), "span-id")
 	defer span.End()
 
 	SetLevel(InfoLevel)
@@ -101,10 +101,10 @@ func TestTraceInfo(t *testing.T) {
 	l.WithDuration(time.Second).Infov(testlog)
 	validate(t, w.String(), true, true)
 	w.Reset()
-	l.WithDuration(time.Second).Infow(testlog, Field("foo", "bar"))
+	l.WithDuration(time.Second).Infow(testlog, Field("basket", "ball"))
 	validate(t, w.String(), true, true)
-	assert.True(t, strings.Contains(w.String(), "foo"), w.String())
-	assert.True(t, strings.Contains(w.String(), "bar"), w.String())
+	assert.True(t, strings.Contains(w.String(), "basket"), w.String())
+	assert.True(t, strings.Contains(w.String(), "ball"), w.String())
 }
 
 func TestTraceInfoConsole(t *testing.T) {
@@ -124,7 +124,7 @@ func TestTraceInfoConsole(t *testing.T) {
 	otel.SetTracerProvider(tp)
 	defer otel.SetTracerProvider(otp)
 
-	ctx, span := tp.Tracer("foo").Start(context.Background(), "bar")
+	ctx, span := tp.Tracer("trace-id").Start(context.Background(), "span-id")
 	defer span.End()
 
 	l := WithContext(ctx)
@@ -153,7 +153,7 @@ func TestTraceSlow(t *testing.T) {
 	otel.SetTracerProvider(tp)
 	defer otel.SetTracerProvider(otp)
 
-	ctx, span := tp.Tracer("foo").Start(context.Background(), "bar")
+	ctx, span := tp.Tracer("trace-id").Start(context.Background(), "span-id")
 	defer span.End()
 
 	l := WithContext(ctx)
@@ -168,10 +168,10 @@ func TestTraceSlow(t *testing.T) {
 	l.WithDuration(time.Second).Slowv(testlog)
 	validate(t, w.String(), true, true)
 	w.Reset()
-	l.WithDuration(time.Second).Sloww(testlog, Field("foo", "bar"))
+	l.WithDuration(time.Second).Sloww(testlog, Field("basket", "ball"))
 	validate(t, w.String(), true, true)
-	assert.True(t, strings.Contains(w.String(), "foo"), w.String())
-	assert.True(t, strings.Contains(w.String(), "bar"), w.String())
+	assert.True(t, strings.Contains(w.String(), "basket"), w.String())
+	assert.True(t, strings.Contains(w.String(), "ball"), w.String())
 }
 
 func TestTraceWithoutContext(t *testing.T) {
@@ -192,6 +192,25 @@ func TestTraceWithoutContext(t *testing.T) {
 	validate(t, w.String(), false, false)
 }
 
+func TestLogWithFields(t *testing.T) {
+	w := new(mockWriter)
+	old := writer.Swap(w)
+	writer.lock.RLock()
+	defer func() {
+		writer.lock.RUnlock()
+		writer.Store(old)
+	}()
+
+	ctx := WithFields(context.Background(), Field("foo", "bar"))
+	l := WithContext(ctx)
+	SetLevel(InfoLevel)
+	l.Info(testlog)
+
+	var val mockValue
+	assert.Nil(t, json.Unmarshal([]byte(w.String()), &val))
+	assert.Equal(t, "bar", val.Foo)
+}
+
 func validate(t *testing.T, body string, expectedTrace, expectedSpan bool) {
 	var val mockValue
 	dec := json.NewDecoder(strings.NewReader(body))
@@ -217,4 +236,5 @@ func validate(t *testing.T, body string, expectedTrace, expectedSpan bool) {
 type mockValue struct {
 	Trace string `json:"trace"`
 	Span  string `json:"span"`
+	Foo   string `json:"foo"`
 }

core/logx/durationlogger.go

@@ -66,7 +66,7 @@ func (l *durationLogger) Sloww(msg string, fields ...LogField) {
 }
 
 func (l *durationLogger) WithContext(ctx context.Context) Logger {
-	return &traceLogger{
+	return &contextLogger{
 		ctx: ctx,
 		logEntry: logEntry{
 			Duration: l.Duration,

core/logx/fields.go

@@ -0,0 +1,18 @@
+package logx
+
+import "context"
+
+var fieldsContextKey contextKey
+
+type contextKey struct{}
+
+// WithFields returns a new context with the given fields.
+func WithFields(ctx context.Context, fields ...LogField) context.Context {
+	if val := ctx.Value(fieldsContextKey); val != nil {
+		if arr, ok := val.([]LogField); ok {
+			return context.WithValue(ctx, fieldsContextKey, append(arr, fields...))
+		}
+	}
+
+	return context.WithValue(ctx, fieldsContextKey, fields)
+}

core/logx/fields_test.go

@@ -0,0 +1,35 @@
+package logx
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestWithFields(t *testing.T) {
+	ctx := WithFields(context.Background(), Field("a", 1), Field("b", 2))
+	vals := ctx.Value(fieldsContextKey)
+	assert.NotNil(t, vals)
+	fields, ok := vals.([]LogField)
+	assert.True(t, ok)
+	assert.EqualValues(t, []LogField{Field("a", 1), Field("b", 2)}, fields)
+}
+
+func TestWithFieldsAppend(t *testing.T) {
+	var dummyKey struct{}
+	ctx := context.WithValue(context.Background(), dummyKey, "dummy")
+	ctx = WithFields(ctx, Field("a", 1), Field("b", 2))
+	ctx = WithFields(ctx, Field("c", 3), Field("d", 4))
+	vals := ctx.Value(fieldsContextKey)
+	assert.NotNil(t, vals)
+	fields, ok := vals.([]LogField)
+	assert.True(t, ok)
+	assert.Equal(t, "dummy", ctx.Value(dummyKey))
+	assert.EqualValues(t, []LogField{
+		Field("a", 1),
+		Field("b", 2),
+		Field("c", 3),
+		Field("d", 4),
+	}, fields)
+}

core/logx/logs.go

@@ -7,6 +7,7 @@ import (
 	"os"
 	"path"
 	"runtime/debug"
+	"sync"
 	"sync/atomic"
 	"time"
 
@@ -20,10 +21,11 @@ var (
 	logLevel   uint32
 	encoding   uint32 = jsonEncodingType
 	// use uint32 for atomic operations
+	disableLog  uint32
 	disableStat uint32
-
-	options logOptions
-	writer  = new(atomicWriter)
+	options     logOptions
+	writer      = new(atomicWriter)
+	setupOnce   sync.Once
 )
 
 type (
@@ -41,6 +43,9 @@ type (
 		gzipEnabled           bool
 		logStackCooldownMills int
 		keepDays              int
+		maxBackups            int
+		maxSize               int
+		rotationRule          string
 	}
 
 	// LogField is a key-value pair that will be added to the log entry.
@@ -69,6 +74,7 @@ func Close() error {
 
 // Disable disables the logging.
 func Disable() {
+	atomic.StoreUint32(&disableLog, 1)
 	writer.Store(nopWriter{})
 }
 
@@ -187,7 +193,6 @@ func MustSetup(c LogConf) {
 
 // Reset clears the writer and resets the log level.
 func Reset() Writer {
-	SetLevel(InfoLevel)
 	return writer.Swap(nil)
 }
 
@@ -197,9 +202,8 @@ func SetLevel(level uint32) {
 }
 
 // SetWriter sets the logging writer. It can be used to customize the logging.
-// Call Reset before calling SetWriter again.
 func SetWriter(w Writer) {
-	if writer.Load() == nil {
+	if atomic.LoadUint32(&disableLog) == 0 {
 		writer.Store(w)
 	}
 }
@@ -207,30 +211,35 @@ func SetWriter(w Writer) {
 // SetUp sets up the logx. If already set up, just return nil.
 // we allow SetUp to be called multiple times, because for example
 // we need to allow different service frameworks to initialize logx respectively.
-// the same logic for SetUp
-func SetUp(c LogConf) error {
-	setupLogLevel(c)
+func SetUp(c LogConf) (err error) {
+	// Just ignore the subsequent SetUp calls.
+	// Because multiple services in one process might call SetUp respectively.
+	// Need to wait for the first caller to complete the execution.
+	setupOnce.Do(func() {
+		setupLogLevel(c)
 
-	if len(c.TimeFormat) > 0 {
-		timeFormat = c.TimeFormat
-	}
+		if len(c.TimeFormat) > 0 {
+			timeFormat = c.TimeFormat
+		}
 
-	switch c.Encoding {
-	case plainEncoding:
-		atomic.StoreUint32(&encoding, plainEncodingType)
-	default:
-		atomic.StoreUint32(&encoding, jsonEncodingType)
-	}
+		switch c.Encoding {
+		case plainEncoding:
+			atomic.StoreUint32(&encoding, plainEncodingType)
+		default:
+			atomic.StoreUint32(&encoding, jsonEncodingType)
+		}
 
-	switch c.Mode {
-	case fileMode:
-		return setupWithFiles(c)
-	case volumeMode:
-		return setupWithVolume(c)
-	default:
-		setupWithConsole()
-		return nil
-	}
+		switch c.Mode {
+		case fileMode:
+			err = setupWithFiles(c)
+		case volumeMode:
+			err = setupWithVolume(c)
+		default:
+			setupWithConsole()
+		}
+	})
+
+	return
 }
 
 // Severe writes v into severe log.
@@ -294,13 +303,40 @@ func WithGzip() LogOption {
 	}
 }
 
+// WithMaxBackups customizes how many log files backups will be kept.
+func WithMaxBackups(count int) LogOption {
+	return func(opts *logOptions) {
+		opts.maxBackups = count
+	}
+}
+
+// WithMaxSize customizes how much space the writing log file can take up.
+func WithMaxSize(size int) LogOption {
+	return func(opts *logOptions) {
+		opts.maxSize = size
+	}
+}
+
+// WithRotation customizes which log rotation rule to use.
+func WithRotation(r string) LogOption {
+	return func(opts *logOptions) {
+		opts.rotationRule = r
+	}
+}
+
 func createOutput(path string) (io.WriteCloser, error) {
 	if len(path) == 0 {
 		return nil, ErrLogPathNotSet
 	}
 
-	return NewLogger(path, DefaultRotateRule(path, backupFileDelimiter, options.keepDays,
-		options.gzipEnabled), options.gzipEnabled)
+	switch options.rotationRule {
+	case sizeRotationRule:
+		return NewLogger(path, NewSizeLimitRotateRule(path, backupFileDelimiter, options.keepDays,
+			options.maxSize, options.maxBackups, options.gzipEnabled), options.gzipEnabled)
+	default:
+		return NewLogger(path, DefaultRotateRule(path, backupFileDelimiter, options.keepDays,
+			options.gzipEnabled), options.gzipEnabled)
+	}
 }
 
 func errorAnySync(v interface{}) {
@@ -324,8 +360,7 @@ func errorTextSync(msg string) {
 func getWriter() Writer {
 	w := writer.Load()
 	if w == nil {
-		w = newConsoleWriter()
-		writer.Store(w)
+		w = writer.StoreIfNil(newConsoleWriter())
 	}
 
 	return w

core/logx/logs_test.go

@@ -4,7 +4,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"io/ioutil"
+	"io"
 	"log"
 	"os"
 	"reflect"
@@ -599,12 +599,14 @@ func TestDisableStat(t *testing.T) {
 }
 
 func TestSetWriter(t *testing.T) {
+	atomic.StoreUint32(&disableLog, 0)
 	Reset()
 	SetWriter(nopWriter{})
 	assert.NotNil(t, writer.Load())
 	assert.True(t, writer.Load() == nopWriter{})
-	SetWriter(new(mockWriter))
-	assert.True(t, writer.Load() == nopWriter{})
+	mocked := new(mockWriter)
+	SetWriter(mocked)
+	assert.Equal(t, mocked, writer.Load())
 }
 
 func TestWithGzip(t *testing.T) {
@@ -647,7 +649,7 @@ func BenchmarkCopyByteSlice(b *testing.B) {
 		buf = make([]byte, len(s))
 		copy(buf, s)
 	}
-	fmt.Fprint(ioutil.Discard, buf)
+	fmt.Fprint(io.Discard, buf)
 }
 
 func BenchmarkCopyOnWriteByteSlice(b *testing.B) {
@@ -656,7 +658,7 @@ func BenchmarkCopyOnWriteByteSlice(b *testing.B) {
 		size := len(s)
 		buf = s[:size:size]
 	}
-	fmt.Fprint(ioutil.Discard, buf)
+	fmt.Fprint(io.Discard, buf)
 }
 
 func BenchmarkCacheByteSlice(b *testing.B) {
@@ -670,7 +672,7 @@ func BenchmarkCacheByteSlice(b *testing.B) {
 func BenchmarkLogs(b *testing.B) {
 	b.ReportAllocs()
 
-	log.SetOutput(ioutil.Discard)
+	log.SetOutput(io.Discard)
 	for i := 0; i < b.N; i++ {
 		Info(i)
 	}

core/logx/readme-cn.md

@@ -8,15 +8,18 @@
 
 ```go
 type LogConf struct {
-	ServiceName         string `json:",optional"`
-	Mode                string `json:",default=console,options=[console,file,volume]"`
-	Encoding            string `json:",default=json,options=[json,plain]"`
-	TimeFormat          string `json:",optional"`
-	Path                string `json:",default=logs"`
-	Level               string `json:",default=info,options=[info,error,severe]"`
-	Compress            bool   `json:",optional"`
-	KeepDays            int    `json:",optional"`
-	StackCooldownMillis int    `json:",default=100"`
+	ServiceName         string              `json:",optional"`
+	Mode                string              `json:",default=console,options=[console,file,volume]"`
+	Encoding            string              `json:",default=json,options=[json,plain]"`
+	TimeFormat          string              `json:",optional"`
+	Path                string              `json:",default=logs"`
+	Level               string              `json:",default=info,options=[info,error,severe]"`
+	Compress            bool                `json:",optional"`
+	KeepDays            int                 `json:",optional"`
+	StackCooldownMillis int                 `json:",default=100"`
+	MaxBackups          int                 `json:",default=0"`
+	MaxSize             int                 `json:",default=0"`
+	Rotation            string              `json:",default=daily,options=[daily,size]"`
 }
 ```
 
@@ -37,6 +40,12 @@ type LogConf struct {
 - `Compress`: 是否压缩日志文件，只在 `file` 模式下工作
 - `KeepDays`：日志文件被保留多少天，在给定的天数之后，过期的文件将被自动删除。对 `console` 模式没有影响
 - `StackCooldownMillis`：多少毫秒后再次写入堆栈跟踪。用来避免堆栈跟踪日志过多
+- `MaxBackups`: 多少个日志文件备份将被保存。0代表所有备份都被保存。当`Rotation`被设置为`size`时才会起作用。注意：`KeepDays`选项的优先级会比`MaxBackups`高，即使`MaxBackups`被设置为0，当达到`KeepDays`上限时备份文件同样会被删除。
+- `MaxSize`: 当前被写入的日志文件最大可占用多少空间。0代表没有上限。单位为`MB`。当`Rotation`被设置为`size`时才会起作用。
+- `Rotation`: 日志轮转策略类型。默认为`daily`（按天轮转）。
+  - `daily` 按天轮转。
+  - `size` 按日志大小轮转。
+
 
 ## 打印日志方法
 

core/logx/readme.md

@@ -8,15 +8,18 @@ English | [简体中文](readme-cn.md)
 
 ```go
 type LogConf struct {
-	ServiceName         string `json:",optional"`
-	Mode                string `json:",default=console,options=[console,file,volume]"`
-	Encoding            string `json:",default=json,options=[json,plain]"`
-	TimeFormat          string `json:",optional"`
-	Path                string `json:",default=logs"`
-	Level               string `json:",default=info,options=[info,error,severe]"`
-	Compress            bool   `json:",optional"`
-	KeepDays            int    `json:",optional"`
-	StackCooldownMillis int    `json:",default=100"`
+	ServiceName         string              `json:",optional"`
+	Mode                string              `json:",default=console,options=[console,file,volume]"`
+	Encoding            string              `json:",default=json,options=[json,plain]"`
+	TimeFormat          string              `json:",optional"`
+	Path                string              `json:",default=logs"`
+	Level               string              `json:",default=info,options=[info,error,severe]"`
+	Compress            bool                `json:",optional"`
+	KeepDays            int                 `json:",optional"`
+	StackCooldownMillis int                 `json:",default=100"`
+	MaxBackups          int                 `json:",default=0"`
+	MaxSize             int                 `json:",default=0"`
+	Rotation            string              `json:",default=daily,options=[daily,size]"`
 }
 ```
 
@@ -37,6 +40,11 @@ type LogConf struct {
 - `Compress`: whether or not to compress log files, only works with `file` mode.
 - `KeepDays`: how many days that the log files are kept, after the given days, the outdated files will be deleted automatically. It has no effect on `console` mode.
 - `StackCooldownMillis`: how many milliseconds to rewrite stacktrace again. It’s used to avoid stacktrace flooding.
+- `MaxBackups`: represents how many backup log files will be kept. 0 means all files will be kept forever. Only take effect when `Rotation` is `size`. NOTE: the level of option `KeepDays` will be higher. Even thougth `MaxBackups` sets 0, log files will still be removed if the `KeepDays` limitation is reached.
+- `MaxSize`: represents how much space the writing log file takes up. 0 means no limit. The unit is `MB`. Only take effect when `Rotation` is `size`.
+- `Rotation`: represents the type of log rotation rule. Default is `daily`.
+  - `daily` rotate the logs by day.
+  - `size` rotate the logs by size of logs.
 
 ## Logging methods
 

core/logx/rotatelogger.go

@@ -9,6 +9,7 @@ import (
 	"os"
 	"path"
 	"path/filepath"
+	"sort"
 	"strings"
 	"sync"
 	"time"
@@ -19,10 +20,13 @@ import (
 
 const (
 	dateFormat      = "2006-01-02"
+	fileTimeFormat  = time.RFC3339
 	hoursPerDay     = 24
 	bufferSize      = 100
 	defaultDirMode  = 0o755
 	defaultFileMode = 0o600
+	gzipExt         = ".gz"
+	megaBytes       = 1 << 20
 )
 
 // ErrLogFileClosed is an error that indicates the log file is already closed.
@@ -34,7 +38,7 @@ type (
 		BackupFileName() string
 		MarkRotated()
 		OutdatedFiles() []string
-		ShallRotate() bool
+		ShallRotate(size int64) bool
 	}
 
 	// A RotateLogger is a Logger that can rotate log files with given rules.
@@ -47,8 +51,9 @@ type (
 		rule     RotateRule
 		compress bool
 		// can't use threading.RoutineGroup because of cycle import
-		waitGroup sync.WaitGroup
-		closeOnce sync.Once
+		waitGroup   sync.WaitGroup
+		closeOnce   sync.Once
+		currentSize int64
 	}
 
 	// A DailyRotateRule is a rule to daily rotate the log files.
@@ -59,6 +64,13 @@ type (
 		days        int
 		gzip        bool
 	}
+
+	// SizeLimitRotateRule a rotation rule that make the log file rotated base on size
+	SizeLimitRotateRule struct {
+		DailyRotateRule
+		maxSize    int64
+		maxBackups int
+	}
 )
 
 // DefaultRotateRule is a default log rotating rule, currently DailyRotateRule.
@@ -90,7 +102,7 @@ func (r *DailyRotateRule) OutdatedFiles() []string {
 
 	var pattern string
 	if r.gzip {
-		pattern = fmt.Sprintf("%s%s*.gz", r.filename, r.delimiter)
+		pattern = fmt.Sprintf("%s%s*%s", r.filename, r.delimiter, gzipExt)
 	} else {
 		pattern = fmt.Sprintf("%s%s*", r.filename, r.delimiter)
 	}
@@ -105,7 +117,7 @@ func (r *DailyRotateRule) OutdatedFiles() []string {
 	boundary := time.Now().Add(-time.Hour * time.Duration(hoursPerDay*r.days)).Format(dateFormat)
 	fmt.Fprintf(&buf, "%s%s%s", r.filename, r.delimiter, boundary)
 	if r.gzip {
-		buf.WriteString(".gz")
+		buf.WriteString(gzipExt)
 	}
 	boundaryFile := buf.String()
 
@@ -120,10 +132,100 @@ func (r *DailyRotateRule) OutdatedFiles() []string {
 }
 
 // ShallRotate checks if the file should be rotated.
-func (r *DailyRotateRule) ShallRotate() bool {
+func (r *DailyRotateRule) ShallRotate(_ int64) bool {
 	return len(r.rotatedTime) > 0 && getNowDate() != r.rotatedTime
 }
 
+// NewSizeLimitRotateRule returns the rotation rule with size limit
+func NewSizeLimitRotateRule(filename, delimiter string, days, maxSize, maxBackups int, gzip bool) RotateRule {
+	return &SizeLimitRotateRule{
+		DailyRotateRule: DailyRotateRule{
+			rotatedTime: getNowDateInRFC3339Format(),
+			filename:    filename,
+			delimiter:   delimiter,
+			days:        days,
+			gzip:        gzip,
+		},
+		maxSize:    int64(maxSize) * megaBytes,
+		maxBackups: maxBackups,
+	}
+}
+
+func (r *SizeLimitRotateRule) BackupFileName() string {
+	dir := filepath.Dir(r.filename)
+	prefix, ext := r.parseFilename()
+	timestamp := getNowDateInRFC3339Format()
+	return filepath.Join(dir, fmt.Sprintf("%s%s%s%s", prefix, r.delimiter, timestamp, ext))
+}
+
+func (r *SizeLimitRotateRule) MarkRotated() {
+	r.rotatedTime = getNowDateInRFC3339Format()
+}
+
+func (r *SizeLimitRotateRule) OutdatedFiles() []string {
+	dir := filepath.Dir(r.filename)
+	prefix, ext := r.parseFilename()
+
+	var pattern string
+	if r.gzip {
+		pattern = fmt.Sprintf("%s%s%s%s*%s%s", dir, string(filepath.Separator),
+			prefix, r.delimiter, ext, gzipExt)
+	} else {
+		pattern = fmt.Sprintf("%s%s%s%s*%s", dir, string(filepath.Separator),
+			prefix, r.delimiter, ext)
+	}
+
+	files, err := filepath.Glob(pattern)
+	if err != nil {
+		Errorf("failed to delete outdated log files, error: %s", err)
+		return nil
+	}
+
+	sort.Strings(files)
+
+	outdated := make(map[string]lang.PlaceholderType)
+
+	// test if too many backups
+	if r.maxBackups > 0 && len(files) > r.maxBackups {
+		for _, f := range files[:len(files)-r.maxBackups] {
+			outdated[f] = lang.Placeholder
+		}
+		files = files[len(files)-r.maxBackups:]
+	}
+
+	// test if any too old backups
+	if r.days > 0 {
+		boundary := time.Now().Add(-time.Hour * time.Duration(hoursPerDay*r.days)).Format(fileTimeFormat)
+		boundaryFile := filepath.Join(dir, fmt.Sprintf("%s%s%s%s", prefix, r.delimiter, boundary, ext))
+		if r.gzip {
+			boundaryFile += gzipExt
+		}
+		for _, f := range files {
+			if f >= boundaryFile {
+				break
+			}
+			outdated[f] = lang.Placeholder
+		}
+	}
+
+	var result []string
+	for k := range outdated {
+		result = append(result, k)
+	}
+	return result
+}
+
+func (r *SizeLimitRotateRule) ShallRotate(size int64) bool {
+	return r.maxSize > 0 && r.maxSize < size
+}
+
+func (r *SizeLimitRotateRule) parseFilename() (prefix, ext string) {
+	logName := filepath.Base(r.filename)
+	ext = filepath.Ext(r.filename)
+	prefix = logName[:len(logName)-len(ext)]
+	return
+}
+
 // NewLogger returns a RotateLogger with given filename and rule, etc.
 func NewLogger(filename string, rule RotateRule, compress bool) (*RotateLogger, error) {
 	l := &RotateLogger{
@@ -282,15 +384,17 @@ func (l *RotateLogger) startWorker() {
 }
 
 func (l *RotateLogger) write(v []byte) {
-	if l.rule.ShallRotate() {
+	if l.rule.ShallRotate(l.currentSize + int64(len(v))) {
 		if err := l.rotate(); err != nil {
 			log.Println(err)
 		} else {
 			l.rule.MarkRotated()
+			l.currentSize = 0
 		}
 	}
 	if l.fp != nil {
 		l.fp.Write(v)
+		l.currentSize += int64(len(v))
 	}
 }
 
@@ -308,6 +412,10 @@ func getNowDate() string {
 	return time.Now().Format(dateFormat)
 }
 
+func getNowDateInRFC3339Format() string {
+	return time.Now().Format(fileTimeFormat)
+}
+
 func gzipFile(file string) error {
 	in, err := os.Open(file)
 	if err != nil {
@@ -315,7 +423,7 @@ func gzipFile(file string) error {
 	}
 	defer in.Close()
 
-	out, err := os.Create(fmt.Sprintf("%s.gz", file))
+	out, err := os.Create(fmt.Sprintf("%s%s", file, gzipExt))
 	if err != nil {
 		return err
 	}

core/logx/rotatelogger_test.go

@@ -29,7 +29,34 @@ func TestDailyRotateRuleOutdatedFiles(t *testing.T) {
 func TestDailyRotateRuleShallRotate(t *testing.T) {
 	var rule DailyRotateRule
 	rule.rotatedTime = time.Now().Add(time.Hour * 24).Format(dateFormat)
-	assert.True(t, rule.ShallRotate())
+	assert.True(t, rule.ShallRotate(0))
+}
+
+func TestSizeLimitRotateRuleMarkRotated(t *testing.T) {
+	var rule SizeLimitRotateRule
+	rule.MarkRotated()
+	assert.Equal(t, getNowDateInRFC3339Format(), rule.rotatedTime)
+}
+
+func TestSizeLimitRotateRuleOutdatedFiles(t *testing.T) {
+	var rule SizeLimitRotateRule
+	assert.Empty(t, rule.OutdatedFiles())
+	rule.days = 1
+	assert.Empty(t, rule.OutdatedFiles())
+	rule.gzip = true
+	assert.Empty(t, rule.OutdatedFiles())
+	rule.maxBackups = 0
+	assert.Empty(t, rule.OutdatedFiles())
+}
+
+func TestSizeLimitRotateRuleShallRotate(t *testing.T) {
+	var rule SizeLimitRotateRule
+	rule.rotatedTime = time.Now().Add(time.Hour * 24).Format(fileTimeFormat)
+	rule.maxSize = 0
+	assert.False(t, rule.ShallRotate(0))
+	rule.maxSize = 100
+	assert.False(t, rule.ShallRotate(0))
+	assert.True(t, rule.ShallRotate(101*megaBytes))
 }
 
 func TestRotateLoggerClose(t *testing.T) {
@@ -142,3 +169,162 @@ func TestRotateLoggerWrite(t *testing.T) {
 func TestLogWriterClose(t *testing.T) {
 	assert.Nil(t, newLogWriter(nil).Close())
 }
+
+func TestRotateLoggerWithSizeLimitRotateRuleClose(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(SizeLimitRotateRule), false)
+	assert.Nil(t, err)
+	assert.Nil(t, logger.Close())
+}
+
+func TestRotateLoggerGetBackupWithSizeLimitRotateRuleFilename(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(SizeLimitRotateRule), false)
+	assert.Nil(t, err)
+	assert.True(t, len(logger.getBackupFilename()) > 0)
+	logger.backup = ""
+	assert.True(t, len(logger.getBackupFilename()) > 0)
+}
+
+func TestRotateLoggerWithSizeLimitRotateRuleMayCompressFile(t *testing.T) {
+	old := os.Stdout
+	os.Stdout = os.NewFile(0, os.DevNull)
+	defer func() {
+		os.Stdout = old
+	}()
+
+	filename, err := fs.TempFilenameWithText("foo")
+	assert.Nil(t, err)
+	if len(filename) > 0 {
+		defer os.Remove(filename)
+	}
+	logger, err := NewLogger(filename, new(SizeLimitRotateRule), false)
+	assert.Nil(t, err)
+	logger.maybeCompressFile(filename)
+	_, err = os.Stat(filename)
+	assert.Nil(t, err)
+}
+
+func TestRotateLoggerWithSizeLimitRotateRuleMayCompressFileTrue(t *testing.T) {
+	old := os.Stdout
+	os.Stdout = os.NewFile(0, os.DevNull)
+	defer func() {
+		os.Stdout = old
+	}()
+
+	filename, err := fs.TempFilenameWithText("foo")
+	assert.Nil(t, err)
+	logger, err := NewLogger(filename, new(SizeLimitRotateRule), true)
+	assert.Nil(t, err)
+	if len(filename) > 0 {
+		defer os.Remove(filepath.Base(logger.getBackupFilename()) + ".gz")
+	}
+	logger.maybeCompressFile(filename)
+	_, err = os.Stat(filename)
+	assert.NotNil(t, err)
+}
+
+func TestRotateLoggerWithSizeLimitRotateRuleRotate(t *testing.T) {
+	filename, err := fs.TempFilenameWithText("foo")
+	assert.Nil(t, err)
+	logger, err := NewLogger(filename, new(SizeLimitRotateRule), true)
+	assert.Nil(t, err)
+	if len(filename) > 0 {
+		defer func() {
+			os.Remove(logger.getBackupFilename())
+			os.Remove(filepath.Base(logger.getBackupFilename()) + ".gz")
+		}()
+	}
+	err = logger.rotate()
+	switch v := err.(type) {
+	case *os.LinkError:
+		// avoid rename error on docker container
+		assert.Equal(t, syscall.EXDEV, v.Err)
+	case *os.PathError:
+		// ignore remove error for tests,
+		// files are cleaned in GitHub actions.
+		assert.Equal(t, "remove", v.Op)
+	default:
+		assert.Nil(t, err)
+	}
+}
+
+func TestRotateLoggerWithSizeLimitRotateRuleWrite(t *testing.T) {
+	filename, err := fs.TempFilenameWithText("foo")
+	assert.Nil(t, err)
+	rule := new(SizeLimitRotateRule)
+	logger, err := NewLogger(filename, rule, true)
+	assert.Nil(t, err)
+	if len(filename) > 0 {
+		defer func() {
+			os.Remove(logger.getBackupFilename())
+			os.Remove(filepath.Base(logger.getBackupFilename()) + ".gz")
+		}()
+	}
+	// the following write calls cannot be changed to Write, because of DATA RACE.
+	logger.write([]byte(`foo`))
+	rule.rotatedTime = time.Now().Add(-time.Hour * 24).Format(dateFormat)
+	logger.write([]byte(`bar`))
+	logger.Close()
+	logger.write([]byte(`baz`))
+}
+
+func BenchmarkRotateLogger(b *testing.B) {
+	filename := "./test.log"
+	filename2 := "./test2.log"
+	dailyRotateRuleLogger, err1 := NewLogger(
+		filename,
+		DefaultRotateRule(
+			filename,
+			backupFileDelimiter,
+			1,
+			true,
+		),
+		true,
+	)
+	if err1 != nil {
+		b.Logf("Failed to new daily rotate rule logger: %v", err1)
+		b.FailNow()
+	}
+	sizeLimitRotateRuleLogger, err2 := NewLogger(
+		filename2,
+		NewSizeLimitRotateRule(
+			filename,
+			backupFileDelimiter,
+			1,
+			100,
+			10,
+			true,
+		),
+		true,
+	)
+	if err2 != nil {
+		b.Logf("Failed to new size limit rotate rule logger: %v", err1)
+		b.FailNow()
+	}
+	defer func() {
+		dailyRotateRuleLogger.Close()
+		sizeLimitRotateRuleLogger.Close()
+		os.Remove(filename)
+		os.Remove(filename2)
+	}()
+
+	b.Run("daily rotate rule", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			dailyRotateRuleLogger.write([]byte("testing\ntesting\n"))
+		}
+	})
+	b.Run("size limit rotate rule", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			sizeLimitRotateRuleLogger.write([]byte("testing\ntesting\n"))
+		}
+	})
+}

core/logx/vars.go

@@ -15,9 +15,9 @@ const (
 	jsonEncodingType = iota
 	plainEncodingType
 
-	jsonEncoding     = "json"
 	plainEncoding    = "plain"
 	plainEncodingSep = '\t'
+	sizeRotationRule = "size"
 )
 
 const (
@@ -27,9 +27,8 @@ const (
 	slowFilename   = "slow.log"
 	statFilename   = "stat.log"
 
-	consoleMode = "console"
-	fileMode    = "file"
-	volumeMode  = "volume"
+	fileMode   = "file"
+	volumeMode = "volume"
 
 	levelAlert  = "alert"
 	levelInfo   = "info"

core/logx/writer.go

@@ -5,12 +5,12 @@ import (
 	"fmt"
 	"io"
 	"log"
-	"os"
 	"path"
 	"strings"
 	"sync"
 	"sync/atomic"
 
+	fatihcolor "github.com/fatih/color"
 	"github.com/zeromicro/go-zero/core/color"
 )
 
@@ -63,21 +63,32 @@ func (w *atomicWriter) Load() Writer {
 
 func (w *atomicWriter) Store(v Writer) {
 	w.lock.Lock()
+	defer w.lock.Unlock()
 	w.writer = v
-	w.lock.Unlock()
+}
+
+func (w *atomicWriter) StoreIfNil(v Writer) Writer {
+	w.lock.Lock()
+	defer w.lock.Unlock()
+
+	if w.writer == nil {
+		w.writer = v
+	}
+
+	return w.writer
 }
 
 func (w *atomicWriter) Swap(v Writer) Writer {
 	w.lock.Lock()
+	defer w.lock.Unlock()
 	old := w.writer
 	w.writer = v
-	w.lock.Unlock()
 	return old
 }
 
 func newConsoleWriter() Writer {
-	outLog := newLogWriter(log.New(os.Stdout, "", flags))
-	errLog := newLogWriter(log.New(os.Stderr, "", flags))
+	outLog := newLogWriter(log.New(fatihcolor.Output, "", flags))
+	errLog := newLogWriter(log.New(fatihcolor.Error, "", flags))
 	return &concreteWriter{
 		infoLog:   outLog,
 		errorLog:  errLog,
@@ -109,6 +120,14 @@ func newFileWriter(c LogConf) (Writer, error) {
 	if c.KeepDays > 0 {
 		opts = append(opts, WithKeepDays(c.KeepDays))
 	}
+	if c.MaxBackups > 0 {
+		opts = append(opts, WithMaxBackups(c.MaxBackups))
+	}
+	if c.MaxSize > 0 {
+		opts = append(opts, WithMaxSize(c.MaxSize))
+	}
+
+	opts = append(opts, WithRotation(c.Rotation))
 
 	accessFile := path.Join(c.Path, accessFilename)
 	errorFile := path.Join(c.Path, errorFilename)
@@ -299,29 +318,7 @@ func writePlainAny(writer io.Writer, level string, val interface{}, fields ...st
 	case fmt.Stringer:
 		writePlainText(writer, level, v.String(), fields...)
 	default:
-		var buf strings.Builder
-		buf.WriteString(getTimestamp())
-		buf.WriteByte(plainEncodingSep)
-		buf.WriteString(level)
-		buf.WriteByte(plainEncodingSep)
-		if err := json.NewEncoder(&buf).Encode(val); err != nil {
-			log.Println(err.Error())
-			return
-		}
-
-		for _, item := range fields {
-			buf.WriteByte(plainEncodingSep)
-			buf.WriteString(item)
-		}
-		buf.WriteByte('\n')
-		if writer == nil {
-			log.Println(buf.String())
-			return
-		}
-
-		if _, err := fmt.Fprint(writer, buf.String()); err != nil {
-			log.Println(err.Error())
-		}
+		writePlainValue(writer, level, v, fields...)
 	}
 }
 
@@ -346,3 +343,29 @@ func writePlainText(writer io.Writer, level, msg string, fields ...string) {
 		log.Println(err.Error())
 	}
 }
+
+func writePlainValue(writer io.Writer, level string, val interface{}, fields ...string) {
+	var buf strings.Builder
+	buf.WriteString(getTimestamp())
+	buf.WriteByte(plainEncodingSep)
+	buf.WriteString(level)
+	buf.WriteByte(plainEncodingSep)
+	if err := json.NewEncoder(&buf).Encode(val); err != nil {
+		log.Println(err.Error())
+		return
+	}
+
+	for _, item := range fields {
+		buf.WriteByte(plainEncodingSep)
+		buf.WriteString(item)
+	}
+	buf.WriteByte('\n')
+	if writer == nil {
+		log.Println(buf.String())
+		return
+	}
+
+	if _, err := fmt.Fprint(writer, buf.String()); err != nil {
+		log.Println(err.Error())
+	}
+}

core/mapping/marshaler_test.go

@@ -261,6 +261,78 @@ func TestMarshal_RangeOut(t *testing.T) {
 	}
 }
 
+func TestMarshal_RangeIllegal(t *testing.T) {
+	tests := []interface{}{
+		struct {
+			Int int `json:"int,range=[3:1]"`
+		}{
+			Int: 2,
+		},
+		struct {
+			Int int `json:"int,range=(3:1]"`
+		}{
+			Int: 2,
+		},
+	}
+
+	for _, test := range tests {
+		_, err := Marshal(test)
+		assert.Equal(t, err, errNumberRange)
+	}
+}
+
+func TestMarshal_RangeLeftEqualsToRight(t *testing.T) {
+	tests := []struct {
+		name  string
+		value interface{}
+		err   error
+	}{
+		{
+			name: "left inclusive, right inclusive",
+			value: struct {
+				Int int `json:"int,range=[2:2]"`
+			}{
+				Int: 2,
+			},
+		},
+		{
+			name: "left inclusive, right exclusive",
+			value: struct {
+				Int int `json:"int,range=[2:2)"`
+			}{
+				Int: 2,
+			},
+			err: errNumberRange,
+		},
+		{
+			name: "left exclusive, right inclusive",
+			value: struct {
+				Int int `json:"int,range=(2:2]"`
+			}{
+				Int: 2,
+			},
+			err: errNumberRange,
+		},
+		{
+			name: "left exclusive, right exclusive",
+			value: struct {
+				Int int `json:"int,range=(2:2)"`
+			}{
+				Int: 2,
+			},
+			err: errNumberRange,
+		},
+	}
+
+	for _, test := range tests {
+		test := test
+		t.Run(test.name, func(t *testing.T) {
+			_, err := Marshal(test.value)
+			assert.Equal(t, test.err, err)
+		})
+	}
+}
+
 func TestMarshal_FromString(t *testing.T) {
 	v := struct {
 		Age int `json:"age,string"`

core/mapping/unmarshaler.go

@@ -534,8 +534,10 @@ func (u *Unmarshaler) fillSliceValue(slice reflect.Value, index int,
 	baseKind reflect.Kind, value interface{}) error {
 	ithVal := slice.Index(index)
 	switch v := value.(type) {
-	case json.Number:
+	case fmt.Stringer:
 		return setValue(baseKind, ithVal, v.String())
+	case string:
+		return setValue(baseKind, ithVal, v)
 	default:
 		// don't need to consider the difference between int, int8, int16, int32, int64,
 		// uint, uint8, uint16, uint32, uint64, because they're handled as json.Number.

core/mapping/unmarshaler_test.go

@@ -2681,7 +2681,7 @@ func TestUnmarshalJsonReaderMultiArray(t *testing.T) {
 	assert.Equal(t, 2, len(res.B))
 }
 
-func TestUnmarshalJsonReaderPtrMultiArray(t *testing.T) {
+func TestUnmarshalJsonReaderPtrMultiArrayString(t *testing.T) {
 	payload := `{"a": "133", "b": [["add", "cccd"], ["eeee"]]}`
 	var res struct {
 		A string      `json:"a"`
@@ -2694,6 +2694,32 @@ func TestUnmarshalJsonReaderPtrMultiArray(t *testing.T) {
 	assert.Equal(t, 2, len(res.B[0]))
 }
 
+func TestUnmarshalJsonReaderPtrMultiArrayString_Int(t *testing.T) {
+	payload := `{"a": "133", "b": [[11, 22], [33]]}`
+	var res struct {
+		A string      `json:"a"`
+		B [][]*string `json:"b"`
+	}
+	reader := strings.NewReader(payload)
+	err := UnmarshalJsonReader(reader, &res)
+	assert.Nil(t, err)
+	assert.Equal(t, 2, len(res.B))
+	assert.Equal(t, 2, len(res.B[0]))
+}
+
+func TestUnmarshalJsonReaderPtrMultiArrayInt(t *testing.T) {
+	payload := `{"a": "133", "b": [[11, 22], [33]]}`
+	var res struct {
+		A string   `json:"a"`
+		B [][]*int `json:"b"`
+	}
+	reader := strings.NewReader(payload)
+	err := UnmarshalJsonReader(reader, &res)
+	assert.Nil(t, err)
+	assert.Equal(t, 2, len(res.B))
+	assert.Equal(t, 2, len(res.B[0]))
+}
+
 func TestUnmarshalJsonReaderPtrArray(t *testing.T) {
 	payload := `{"a": "133", "b": ["add", "cccd", "eeee"]}`
 	var res struct {
@@ -2706,6 +2732,30 @@ func TestUnmarshalJsonReaderPtrArray(t *testing.T) {
 	assert.Equal(t, 3, len(res.B))
 }
 
+func TestUnmarshalJsonReaderPtrArray_Int(t *testing.T) {
+	payload := `{"a": "133", "b": [11, 22, 33]}`
+	var res struct {
+		A string    `json:"a"`
+		B []*string `json:"b"`
+	}
+	reader := strings.NewReader(payload)
+	err := UnmarshalJsonReader(reader, &res)
+	assert.Nil(t, err)
+	assert.Equal(t, 3, len(res.B))
+}
+
+func TestUnmarshalJsonReaderPtrInt(t *testing.T) {
+	payload := `{"a": "133", "b": [11, 22, 33]}`
+	var res struct {
+		A string    `json:"a"`
+		B []*string `json:"b"`
+	}
+	reader := strings.NewReader(payload)
+	err := UnmarshalJsonReader(reader, &res)
+	assert.Nil(t, err)
+	assert.Equal(t, 3, len(res.B))
+}
+
 func TestUnmarshalJsonWithoutKey(t *testing.T) {
 	payload := `{"A": "1", "B": "2"}`
 	var res struct {

core/mapping/utils.go

@@ -143,6 +143,23 @@ func doParseKeyAndOptions(field reflect.StructField, value string) (string, *fie
 	return key, &fieldOpts, nil
 }
 
+// ensureValue ensures nested members not to be nil.
+// If pointer value is nil, set to a new value.
+func ensureValue(v reflect.Value) reflect.Value {
+	for {
+		if v.Kind() != reflect.Ptr {
+			break
+		}
+
+		if v.IsNil() {
+			v.Set(reflect.New(v.Type().Elem()))
+		}
+		v = v.Elem()
+	}
+
+	return v
+}
+
 func implicitValueRequiredStruct(tag string, tp reflect.Type) (bool, error) {
 	numFields := tp.NumField()
 	for i := 0; i < numFields; i++ {
@@ -294,6 +311,20 @@ func parseNumberRange(str string) (*numberRange, error) {
 		right = math.MaxFloat64
 	}
 
+	if left > right {
+		return nil, errNumberRange
+	}
+
+	// [2:2] valid
+	// [2:2) invalid
+	// (2:2] invalid
+	// (2:2) invalid
+	if left == right {
+		if !leftInclude || !rightInclude {
+			return nil, errNumberRange
+		}
+	}
+
 	return &numberRange{
 		left:         left,
 		leftInclude:  leftInclude,
@@ -478,6 +509,7 @@ func setValue(kind reflect.Kind, value reflect.Value, str string) error {
 		return errValueNotSettable
 	}
 
+	value = ensureValue(value)
 	v, err := convertType(kind, str)
 	if err != nil {
 		return err

core/mapping/utils_test.go

@@ -15,7 +15,7 @@ type Foo struct {
 	StrWithTagAndOption string `key:"stringwithtag,string"`
 }
 
-func TestDeferInt(t *testing.T) {
+func TestDerefInt(t *testing.T) {
 	i := 1
 	s := "hello"
 	number := struct {
@@ -60,6 +60,51 @@ func TestDeferInt(t *testing.T) {
 	}
 }
 
+func TestDerefValInt(t *testing.T) {
+	i := 1
+	s := "hello"
+	number := struct {
+		f float64
+	}{
+		f: 6.4,
+	}
+	cases := []struct {
+		t      reflect.Value
+		expect reflect.Kind
+	}{
+		{
+			t:      reflect.ValueOf(i),
+			expect: reflect.Int,
+		},
+		{
+			t:      reflect.ValueOf(&i),
+			expect: reflect.Int,
+		},
+		{
+			t:      reflect.ValueOf(s),
+			expect: reflect.String,
+		},
+		{
+			t:      reflect.ValueOf(&s),
+			expect: reflect.String,
+		},
+		{
+			t:      reflect.ValueOf(number.f),
+			expect: reflect.Float64,
+		},
+		{
+			t:      reflect.ValueOf(&number.f),
+			expect: reflect.Float64,
+		},
+	}
+
+	for _, each := range cases {
+		t.Run(each.t.String(), func(t *testing.T) {
+			assert.Equal(t, each.expect, ensureValue(each.t).Kind())
+		})
+	}
+}
+
 func TestParseKeyAndOptionWithoutTag(t *testing.T) {
 	var foo Foo
 	rte := reflect.TypeOf(&foo).Elem()

core/metric/counter.go

@@ -3,6 +3,7 @@ package metric
 import (
 	prom "github.com/prometheus/client_golang/prometheus"
 	"github.com/zeromicro/go-zero/core/proc"
+	"github.com/zeromicro/go-zero/core/prometheus"
 )
 
 type (
@@ -47,10 +48,18 @@ func NewCounterVec(cfg *CounterVecOpts) CounterVec {
 }
 
 func (cv *promCounterVec) Inc(labels ...string) {
+	if !prometheus.Enabled() {
+		return
+	}
+
 	cv.counter.WithLabelValues(labels...).Inc()
 }
 
 func (cv *promCounterVec) Add(v float64, labels ...string) {
+	if !prometheus.Enabled() {
+		return
+	}
+
 	cv.counter.WithLabelValues(labels...).Add(v)
 }
 

core/metric/counter_test.go

@@ -5,6 +5,7 @@ import (
 
 	"github.com/prometheus/client_golang/prometheus/testutil"
 	"github.com/stretchr/testify/assert"
+	"github.com/zeromicro/go-zero/core/prometheus"
 )
 
 func TestNewCounterVec(t *testing.T) {
@@ -21,6 +22,7 @@ func TestNewCounterVec(t *testing.T) {
 }
 
 func TestCounterIncr(t *testing.T) {
+	startAgent()
 	counterVec := NewCounterVec(&CounterVecOpts{
 		Namespace: "http_client",
 		Subsystem: "call",
@@ -37,6 +39,7 @@ func TestCounterIncr(t *testing.T) {
 }
 
 func TestCounterAdd(t *testing.T) {
+	startAgent()
 	counterVec := NewCounterVec(&CounterVecOpts{
 		Namespace: "rpc_server",
 		Subsystem: "requests",
@@ -51,3 +54,11 @@ func TestCounterAdd(t *testing.T) {
 	r := testutil.ToFloat64(cv.counter)
 	assert.Equal(t, float64(33), r)
 }
+
+func startAgent() {
+	prometheus.StartAgent(prometheus.Config{
+		Host: "127.0.0.1",
+		Port: 9101,
+		Path: "/metrics",
+	})
+}

core/metric/gauge.go

@@ -3,6 +3,7 @@ package metric
 import (
 	prom "github.com/prometheus/client_golang/prometheus"
 	"github.com/zeromicro/go-zero/core/proc"
+	"github.com/zeromicro/go-zero/core/prometheus"
 )
 
 type (
@@ -50,14 +51,26 @@ func NewGaugeVec(cfg *GaugeVecOpts) GaugeVec {
 }
 
 func (gv *promGaugeVec) Inc(labels ...string) {
+	if !prometheus.Enabled() {
+		return
+	}
+
 	gv.gauge.WithLabelValues(labels...).Inc()
 }
 
 func (gv *promGaugeVec) Add(v float64, labels ...string) {
+	if !prometheus.Enabled() {
+		return
+	}
+
 	gv.gauge.WithLabelValues(labels...).Add(v)
 }
 
 func (gv *promGaugeVec) Set(v float64, labels ...string) {
+	if !prometheus.Enabled() {
+		return
+	}
+
 	gv.gauge.WithLabelValues(labels...).Set(v)
 }
 

core/metric/gauge_test.go

@@ -21,6 +21,7 @@ func TestNewGaugeVec(t *testing.T) {
 }
 
 func TestGaugeInc(t *testing.T) {
+	startAgent()
 	gaugeVec := NewGaugeVec(&GaugeVecOpts{
 		Namespace: "rpc_client2",
 		Subsystem: "requests",
@@ -37,6 +38,7 @@ func TestGaugeInc(t *testing.T) {
 }
 
 func TestGaugeAdd(t *testing.T) {
+	startAgent()
 	gaugeVec := NewGaugeVec(&GaugeVecOpts{
 		Namespace: "rpc_client",
 		Subsystem: "request",
@@ -53,6 +55,7 @@ func TestGaugeAdd(t *testing.T) {
 }
 
 func TestGaugeSet(t *testing.T) {
+	startAgent()
 	gaugeVec := NewGaugeVec(&GaugeVecOpts{
 		Namespace: "http_client",
 		Subsystem: "request",

core/metric/histogram.go

@@ -3,6 +3,7 @@ package metric
 import (
 	prom "github.com/prometheus/client_golang/prometheus"
 	"github.com/zeromicro/go-zero/core/proc"
+	"github.com/zeromicro/go-zero/core/prometheus"
 )
 
 type (
@@ -53,6 +54,10 @@ func NewHistogramVec(cfg *HistogramVecOpts) HistogramVec {
 }
 
 func (hv *promHistogramVec) Observe(v int64, labels ...string) {
+	if !prometheus.Enabled() {
+		return
+	}
+
 	hv.histogram.WithLabelValues(labels...).Observe(float64(v))
 }
 

core/metric/histogram_test.go

@@ -21,6 +21,7 @@ func TestNewHistogramVec(t *testing.T) {
 }
 
 func TestHistogramObserve(t *testing.T) {
+	startAgent()
 	histogramVec := NewHistogramVec(&HistogramVecOpts{
 		Name:    "counts",
 		Help:    "rpc server requests duration(ms).",

core/mr/mapreduce.go

@@ -212,6 +212,8 @@ func mapReduceWithPanicChan(source <-chan interface{}, panicChan *onceChan, mapp
 		cancel(context.DeadlineExceeded)
 		return nil, context.DeadlineExceeded
 	case v := <-panicChan.channel:
+		// drain output here, otherwise for loop panic in defer
+		drain(output)
 		panic(v)
 	case v, ok := <-output:
 		if err := retErr.Load(); err != nil {
@@ -376,9 +378,7 @@ type onceChan struct {
 }
 
 func (oc *onceChan) write(val interface{}) {
-	if atomic.AddInt32(&oc.wrote, 1) > 1 {
-		return
+	if atomic.CompareAndSwapInt32(&oc.wrote, 0, 1) {
+		oc.channel <- val
 	}
-
-	oc.channel <- val
 }

core/mr/mapreduce_test.go

@@ -3,7 +3,7 @@ package mr
 import (
 	"context"
 	"errors"
-	"io/ioutil"
+	"io"
 	"log"
 	"runtime"
 	"sync/atomic"
@@ -17,7 +17,7 @@ import (
 var errDummy = errors.New("dummy")
 
 func init() {
-	log.SetOutput(ioutil.Discard)
+	log.SetOutput(io.Discard)
 }
 
 func TestFinish(t *testing.T) {

core/service/serviceconf.go

@@ -5,6 +5,7 @@ import (
 
 	"github.com/zeromicro/go-zero/core/load"
 	"github.com/zeromicro/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/core/proc"
 	"github.com/zeromicro/go-zero/core/prometheus"
 	"github.com/zeromicro/go-zero/core/stat"
 	"github.com/zeromicro/go-zero/core/trace"
@@ -56,6 +57,9 @@ func (sc ServiceConf) SetUp() error {
 		sc.Telemetry.Name = sc.Name
 	}
 	trace.StartAgent(sc.Telemetry)
+	proc.AddShutdownListener(func() {
+		trace.StopAgent()
+	})
 
 	if len(sc.MetricsUrl) > 0 {
 		stat.SetReportWriter(stat.NewRemoteWriter(sc.MetricsUrl))

core/stat/internal/cpu_linux.go

@@ -4,6 +4,7 @@ import (
 	"errors"
 	"fmt"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/zeromicro/go-zero/core/iox"
@@ -20,10 +21,11 @@ var (
 	preTotal  uint64
 	quota     float64
 	cores     uint64
+	initOnce  sync.Once
 )
 
 // if /proc not present, ignore the cpu calculation, like wsl linux
-func init() {
+func initialize() {
 	cpus, err := cpuSets()
 	if err != nil {
 		logx.Error(err)
@@ -69,10 +71,13 @@ func init() {
 
 // RefreshCpu refreshes cpu usage and returns.
 func RefreshCpu() uint64 {
+	initOnce.Do(initialize)
+
 	total, err := totalCpuUsage()
 	if err != nil {
 		return 0
 	}
+
 	system, err := systemCpuUsage()
 	if err != nil {
 		return 0

core/stores/cache/cachenode.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"math"
 	"math/rand"
 	"sync"
 	"time"
@@ -130,7 +131,7 @@ func (c cacheNode) SetWithExpireCtx(ctx context.Context, key string, val interfa
 		return err
 	}
 
-	return c.rds.SetexCtx(ctx, key, string(data), int(expire.Seconds()))
+	return c.rds.SetexCtx(ctx, key, string(data), int(math.Ceil(expire.Seconds())))
 }
 
 // String returns a string that represents the cacheNode.
@@ -275,5 +276,6 @@ func (c cacheNode) processCache(ctx context.Context, key, data string, v interfa
 }
 
 func (c cacheNode) setCacheWithNotFound(ctx context.Context, key string) error {
-	return c.rds.SetexCtx(ctx, key, notFoundPlaceholder, int(c.aroundDuration(c.notFoundExpiry).Seconds()))
+	seconds := int(math.Ceil(c.aroundDuration(c.notFoundExpiry).Seconds()))
+	return c.rds.SetexCtx(ctx, key, notFoundPlaceholder, seconds)
 }

core/stores/kv/store.go

@@ -110,7 +110,9 @@ type (
 		Ttl(key string) (int, error)
 		TtlCtx(ctx context.Context, key string) (int, error)
 		Zadd(key string, score int64, value string) (bool, error)
+		ZaddFloat(key string, score float64, value string) (bool, error)
 		ZaddCtx(ctx context.Context, key string, score int64, value string) (bool, error)
+		ZaddFloatCtx(ctx context.Context, key string, score float64, value string) (bool, error)
 		Zadds(key string, ps ...redis.Pair) (int64, error)
 		ZaddsCtx(ctx context.Context, key string, ps ...redis.Pair) (int64, error)
 		Zcard(key string) (int, error)
@@ -787,13 +789,21 @@ func (cs clusterStore) Zadd(key string, score int64, value string) (bool, error)
 	return cs.ZaddCtx(context.Background(), key, score, value)
 }
 
+func (cs clusterStore) ZaddFloat(key string, score float64, value string) (bool, error) {
+	return cs.ZaddFloatCtx(context.Background(), key, score, value)
+}
+
 func (cs clusterStore) ZaddCtx(ctx context.Context, key string, score int64, value string) (bool, error) {
+	return cs.ZaddFloatCtx(ctx, key, float64(score), value)
+}
+
+func (cs clusterStore) ZaddFloatCtx(ctx context.Context, key string, score float64, value string) (bool, error) {
 	node, err := cs.getRedis(key)
 	if err != nil {
 		return false, err
 	}
 
-	return node.ZaddCtx(ctx, key, score, value)
+	return node.ZaddFloatCtx(ctx, key, score, value)
 }
 
 func (cs clusterStore) Zadds(key string, ps ...redis.Pair) (int64, error) {

core/stores/kv/store_test.go

@@ -22,7 +22,7 @@ func TestRedis_Decr(t *testing.T) {
 	_, err := store.Decr("a")
 	assert.NotNil(t, err)
 
-	runOnCluster(t, func(client Store) {
+	runOnCluster(func(client Store) {
 		val, err := client.Decr("a")
 		assert.Nil(t, err)
 		assert.Equal(t, int64(-1), val)
@@ -37,7 +37,7 @@ func TestRedis_DecrBy(t *testing.T) {
 	_, err := store.Incrby("a", 2)
 	assert.NotNil(t, err)
 
-	runOnCluster(t, func(client Store) {
+	runOnCluster(func(client Store) {
 		val, err := client.Decrby("a", 2)
 		assert.Nil(t, err)
 		assert.Equal(t, int64(-2), val)
@@ -52,7 +52,7 @@ func TestRedis_Exists(t *testing.T) {
 	_, err := store.Exists("foo")
 	assert.NotNil(t, err)
 
-	runOnCluster(t, func(client Store) {
+	runOnCluster(func(client Store) {
 		ok, err := client.Exists("a")
 		assert.Nil(t, err)
 		assert.False(t, ok)
@@ -68,7 +68,7 @@ func TestRedis_Eval(t *testing.T) {
 	_, err := store.Eval(`redis.call("EXISTS", KEYS[1])`, "key1")
 	assert.NotNil(t, err)
 
-	runOnCluster(t, func(client Store) {
+	runOnCluster(func(client Store) {
 		_, err := client.Eval(`redis.call("EXISTS", KEYS[1])`, "notexist")
 		assert.Equal(t, redis.Nil, err)
 		err = client.Set("key1", "value1")
@@ -88,7 +88,7 @@ func TestRedis_Hgetall(t *testing.T) {
 	_, err = store.Hgetall("a")
 	assert.NotNil(t, err)
 
-	runOnCluster(t, func(client Store) {
+	runOnCluster(func(client Store) {
 		assert.Nil(t, client.Hset("a", "aa", "aaa"))
 		assert.Nil(t, client.Hset("a", "bb", "bbb"))
 		vals, err := client.Hgetall("a")
@@ -105,7 +105,7 @@ func TestRedis_Hvals(t *testing.T) {
 	_, err := store.Hvals("a")
 	assert.NotNil(t, err)
 
-	runOnCluster(t, func(client Store) {
+	runOnCluster(func(client Store) {
 		assert.Nil(t, client.Hset("a", "aa", "aaa"))
 		assert.Nil(t, client.Hset("a", "bb", "bbb"))
 		vals, err := client.Hvals("a")
@@ -119,7 +119,7 @@ func TestRedis_Hsetnx(t *testing.T) {
 	_, err := store.Hsetnx("a", "dd", "ddd")
 	assert.NotNil(t, err)
 
-	runOnCluster(t, func(client Store) {
+	runOnCluster(func(client Store) {
 		assert.Nil(t, client.Hset("a", "aa", "aaa"))
 		assert.Nil(t, client.Hset("a", "bb", "bbb"))
 		ok, err := client.Hsetnx("a", "bb", "ccc")
@@ -141,7 +141,7 @@ func TestRedis_HdelHlen(t *testing.T) {
 	_, err = store.Hlen("a")
 	assert.NotNil(t, err)
 
-	runOnCluster(t, func(client Store) {
+	runOnCluster(func(client Store) {
 		assert.Nil(t, client.Hset("a", "aa", "aaa"))
 		assert.Nil(t, client.Hset("a", "bb", "bbb"))
 		num, err := client.Hlen("a")
@@ -161,7 +161,7 @@ func TestRedis_HIncrBy(t *testing.T) {
 	_, err := store.Hincrby("key", "field", 3)
 	assert.NotNil(t, err)
 
-	runOnCluster(t, func(client Store) {
+	runOnCluster(func(client Store) {
 		val, err := client.Hincrby("key", "field", 2)
 		assert.Nil(t, err)
 		assert.Equal(t, 2, val)
@@ -176,7 +176,7 @@ func TestRedis_Hkeys(t *testing.T) {
 	_, err := store.Hkeys("a")
 	assert.NotNil(t, err)
 
-	runOnCluster(t, func(client Store) {
+	runOnCluster(func(client Store) {
 		assert.Nil(t, client.Hset("a", "aa", "aaa"))
 		assert.Nil(t, client.Hset("a", "bb", "bbb"))
 		vals, err := client.Hkeys("a")
@@ -190,7 +190,7 @@ func TestRedis_Hmget(t *testing.T) {
 	_, err := store.Hmget("a", "aa", "bb")
 	assert.NotNil(t, err)
 
-	runOnCluster(t, func(client Store) {
+	runOnCluster(func(client Store) {
 		assert.Nil(t, client.Hset("a", "aa", "aaa"))
 		assert.Nil(t, client.Hset("a", "bb", "bbb"))
 		vals, err := client.Hmget("a", "aa", "bb")
@@ -209,7 +209,7 @@ func TestRedis_Hmset(t *testing.T) {
 	})
 	assert.NotNil(t, err)
 
-	runOnCluster(t, func(client Store) {
+	runOnCluster(func(client Store) {
 		assert.Nil(t, client.Hmset("a", map[string]string{
 			"aa": "aaa",
 			"bb": "bbb",
@@ -225,7 +225,7 @@ func TestRedis_Incr(t *testing.T) {
 	_, err := store.Incr("a")
 	assert.NotNil(t, err)
 
-	runOnCluster(t, func(client Store) {
+	runOnCluster(func(client Store) {
 		val, err := client.Incr("a")
 		assert.Nil(t, err)
 		assert.Equal(t, int64(1), val)
@@ -240,7 +240,7 @@ func TestRedis_IncrBy(t *testing.T) {
 	_, err := store.Incrby("a", 2)
 	assert.NotNil(t, err)
 
-	runOnCluster(t, func(client Store) {
+	runOnCluster(func(client Store) {
 		val, err := client.Incrby("a", 2)
 		assert.Nil(t, err)
 		assert.Equal(t, int64(2), val)
@@ -267,7 +267,7 @@ func TestRedis_List(t *testing.T) {
 	_, err = store.Lindex("key", 0)
 	assert.NotNil(t, err)
 
-	runOnCluster(t, func(client Store) {
+	runOnCluster(func(client Store) {
 		val, err := client.Lpush("key", "value1", "value2")
 		assert.Nil(t, err)
 		assert.Equal(t, 2, val)
@@ -316,7 +316,7 @@ func TestRedis_Persist(t *testing.T) {
 	err = store.Expireat("key", time.Now().Unix()+5)
 	assert.NotNil(t, err)
 
-	runOnCluster(t, func(client Store) {
+	runOnCluster(func(client Store) {
 		ok, err := client.Persist("key")
 		assert.Nil(t, err)
 		assert.False(t, ok)
@@ -348,7 +348,7 @@ func TestRedis_Sscan(t *testing.T) {
 	_, err = store.Del(key)
 	assert.NotNil(t, err)
 
-	runOnCluster(t, func(client Store) {
+	runOnCluster(func(client Store) {
 		var list []string
 		for i := 0; i < 1550; i++ {
 			list = append(list, stringx.Randn(i))
@@ -390,7 +390,7 @@ func TestRedis_Set(t *testing.T) {
 	_, err = store.Spop("key")
 	assert.NotNil(t, err)
 
-	runOnCluster(t, func(client Store) {
+	runOnCluster(func(client Store) {
 		num, err := client.Sadd("key", 1, 2, 3, 4)
 		assert.Nil(t, err)
 		assert.Equal(t, 4, num)
@@ -434,7 +434,7 @@ func TestRedis_SetGetDel(t *testing.T) {
 	_, err = store.Del("hello")
 	assert.NotNil(t, err)
 
-	runOnCluster(t, func(client Store) {
+	runOnCluster(func(client Store) {
 		err := client.Set("hello", "world")
 		assert.Nil(t, err)
 		val, err := client.Get("hello")
@@ -457,7 +457,7 @@ func TestRedis_SetExNx(t *testing.T) {
 	_, err = store.SetnxEx("newhello", "newworld", 5)
 	assert.NotNil(t, err)
 
-	runOnCluster(t, func(client Store) {
+	runOnCluster(func(client Store) {
 		err := client.Setex("hello", "world", 5)
 		assert.Nil(t, err)
 		ok, err := client.Setnx("hello", "newworld")
@@ -495,7 +495,7 @@ func TestRedis_Getset(t *testing.T) {
 	_, err := store.GetSet("hello", "world")
 	assert.NotNil(t, err)
 
-	runOnCluster(t, func(client Store) {
+	runOnCluster(func(client Store) {
 		val, err := client.GetSet("hello", "world")
 		assert.Nil(t, err)
 		assert.Equal(t, "", val)
@@ -524,7 +524,7 @@ func TestRedis_SetGetDelHashField(t *testing.T) {
 	_, err = store.Hdel("key", "field")
 	assert.NotNil(t, err)
 
-	runOnCluster(t, func(client Store) {
+	runOnCluster(func(client Store) {
 		err := client.Hset("key", "field", "value")
 		assert.Nil(t, err)
 		val, err := client.Hget("key", "field")
@@ -587,8 +587,8 @@ func TestRedis_SortedSet(t *testing.T) {
 	})
 	assert.NotNil(t, err)
 
-	runOnCluster(t, func(client Store) {
-		ok, err := client.Zadd("key", 1, "value1")
+	runOnCluster(func(client Store) {
+		ok, err := client.ZaddFloat("key", 1, "value1")
 		assert.Nil(t, err)
 		assert.True(t, ok)
 		ok, err = client.Zadd("key", 2, "value1")
@@ -724,7 +724,7 @@ func TestRedis_HyperLogLog(t *testing.T) {
 	_, err = store.Pfcount("key")
 	assert.NotNil(t, err)
 
-	runOnCluster(t, func(cluster Store) {
+	runOnCluster(func(cluster Store) {
 		ok, err := cluster.Pfadd("key", "value")
 		assert.Nil(t, err)
 		assert.True(t, ok)
@@ -734,7 +734,7 @@ func TestRedis_HyperLogLog(t *testing.T) {
 	})
 }
 
-func runOnCluster(t *testing.T, fn func(cluster Store)) {
+func runOnCluster(fn func(cluster Store)) {
 	s1.FlushAll()
 	s2.FlushAll()
 

core/stores/mon/collection_test.go

@@ -20,10 +20,6 @@ import (
 
 var errDummy = errors.New("dummy")
 
-func init() {
-	logx.Disable()
-}
-
 func TestKeepPromise_accept(t *testing.T) {
 	p := new(mockPromise)
 	kp := keepablePromise{

core/stores/mongo/collection_test.go

@@ -17,10 +17,6 @@ import (
 
 var errDummy = errors.New("dummy")
 
-func init() {
-	logx.Disable()
-}
-
 func TestKeepPromise_accept(t *testing.T) {
 	p := new(mockPromise)
 	kp := keepablePromise{

core/stores/mongoc/cachedcollection_test.go

@@ -3,7 +3,7 @@ package mongoc
 import (
 	"encoding/json"
 	"errors"
-	"io/ioutil"
+	"io"
 	"log"
 	"os"
 	"runtime"
@@ -117,7 +117,7 @@ func TestStat(t *testing.T) {
 
 func TestStatCacheFails(t *testing.T) {
 	resetStats()
-	log.SetOutput(ioutil.Discard)
+	log.SetOutput(io.Discard)
 	defer log.SetOutput(os.Stdout)
 
 	r := redis.New("localhost:59999")

core/stores/redis/hook.go

@@ -56,6 +56,11 @@ func (h hook) AfterProcess(ctx context.Context, cmd red.Cmder) error {
 		logDuration(ctx, cmd, duration)
 	}
 
+	metricReqDur.Observe(int64(duration/time.Millisecond), cmd.Name())
+	if msg := formatError(err); len(msg) > 0 {
+		metricReqErr.Inc(cmd.Name(), msg)
+	}
+
 	return nil
 }
 
@@ -98,9 +103,43 @@ func (h hook) AfterProcessPipeline(ctx context.Context, cmds []red.Cmder) error
 		logDuration(ctx, cmds[0], duration)
 	}
 
+	metricReqDur.Observe(int64(duration/time.Millisecond), "Pipeline")
+	if msg := formatError(batchError.Err()); len(msg) > 0 {
+		metricReqErr.Inc("Pipeline", msg)
+	}
+
 	return nil
 }
 
+func formatError(err error) string {
+	if err == nil || err == red.Nil {
+		return ""
+	}
+
+	es := err.Error()
+	switch {
+	case strings.HasPrefix(es, "read"):
+		return "read timeout"
+	case strings.HasPrefix(es, "dial"):
+		if strings.Contains(es, "connection refused") {
+			return "connection refused"
+		}
+		return "dial timeout"
+	case strings.HasPrefix(es, "write"):
+		return "write timeout"
+	case strings.Contains(es, "EOF"):
+		return "eof"
+	case strings.Contains(es, "reset"):
+		return "reset"
+	case strings.Contains(es, "broken"):
+		return "broken pipe"
+	case strings.Contains(es, "breaker"):
+		return "breaker"
+	default:
+		return "unexpected error"
+	}
+}
+
 func logDuration(ctx context.Context, cmd red.Cmder, duration time.Duration) {
 	var buf strings.Builder
 	for i, arg := range cmd.Args() {

core/stores/redis/hook_test.go

@@ -21,6 +21,7 @@ func TestHookProcessCase1(t *testing.T) {
 		Batcher:  "jaeger",
 		Sampler:  1.0,
 	})
+	defer ztrace.StopAgent()
 
 	writer := log.Writer()
 	var buf strings.Builder
@@ -44,6 +45,7 @@ func TestHookProcessCase2(t *testing.T) {
 		Batcher:  "jaeger",
 		Sampler:  1.0,
 	})
+	defer ztrace.StopAgent()
 
 	w, restore := injectLog()
 	defer restore()
@@ -108,6 +110,7 @@ func TestHookProcessPipelineCase2(t *testing.T) {
 		Batcher:  "jaeger",
 		Sampler:  1.0,
 	})
+	defer ztrace.StopAgent()
 
 	w, restore := injectLog()
 	defer restore()

core/stores/redis/metrics.go

@@ -0,0 +1,23 @@
+package redis
+
+import "github.com/zeromicro/go-zero/core/metric"
+
+const namespace = "redis_client"
+
+var (
+	metricReqDur = metric.NewHistogramVec(&metric.HistogramVecOpts{
+		Namespace: namespace,
+		Subsystem: "requests",
+		Name:      "duration_ms",
+		Help:      "redis client requests duration(ms).",
+		Labels:    []string{"command"},
+		Buckets:   []float64{5, 10, 25, 50, 100, 250, 500, 1000, 2500},
+	})
+	metricReqErr = metric.NewCounterVec(&metric.CounterVecOpts{
+		Namespace: namespace,
+		Subsystem: "requests",
+		Name:      "error_total",
+		Help:      "redis client requests error count.",
+		Labels:    []string{"command", "error"},
+	})
+)

core/stores/redis/redis.go

@@ -1836,8 +1836,19 @@ func (s *Redis) Zadd(key string, score int64, value string) (bool, error) {
 	return s.ZaddCtx(context.Background(), key, score, value)
 }
 
+// ZaddFloat is the implementation of redis zadd command.
+func (s *Redis) ZaddFloat(key string, score float64, value string) (bool, error) {
+	return s.ZaddFloatCtx(context.Background(), key, score, value)
+}
+
 // ZaddCtx is the implementation of redis zadd command.
 func (s *Redis) ZaddCtx(ctx context.Context, key string, score int64, value string) (
+	val bool, err error) {
+	return s.ZaddFloatCtx(ctx, key, float64(score), value)
+}
+
+// ZaddFloatCtx is the implementation of redis zadd command.
+func (s *Redis) ZaddFloatCtx(ctx context.Context, key string, score float64, value string) (
 	val bool, err error) {
 	err = s.brk.DoWithAcceptable(func() error {
 		conn, err := getRedis(s)
@@ -1846,7 +1857,7 @@ func (s *Redis) ZaddCtx(ctx context.Context, key string, score int64, value stri
 		}
 
 		v, err := conn.ZAdd(ctx, key, &red.Z{
-			Score:  float64(score),
+			Score:  score,
 			Member: value,
 		}).Result()
 		if err != nil {

core/stores/redis/redis_test.go

@@ -810,7 +810,7 @@ func TestRedis_SetGetDelHashField(t *testing.T) {
 
 func TestRedis_SortedSet(t *testing.T) {
 	runOnRedis(t, func(client *Redis) {
-		ok, err := client.Zadd("key", 1, "value1")
+		ok, err := client.ZaddFloat("key", 1, "value1")
 		assert.Nil(t, err)
 		assert.True(t, ok)
 		ok, err = client.Zadd("key", 2, "value1")
@@ -988,8 +988,8 @@ func TestRedis_SortedSet(t *testing.T) {
 		assert.Equal(t, 0, len(pairs))
 		_, err = New(client.Addr, badType()).Zrevrank("key", "value")
 		assert.NotNil(t, err)
-		client.Zadd("second", 2, "aa")
-		client.Zadd("third", 3, "bbb")
+		_, _ = client.Zadd("second", 2, "aa")
+		_, _ = client.Zadd("third", 3, "bbb")
 		val, err = client.Zunionstore("union", &ZStore{
 			Keys:      []string{"second", "third"},
 			Weights:   []float64{1, 2},
@@ -1176,7 +1176,7 @@ func runOnRedis(t *testing.T, fn func(client *Redis)) {
 		}
 
 		if client != nil {
-			client.Close()
+			_ = client.Close()
 		}
 	}()
 	fn(New(s.Addr()))
@@ -1198,7 +1198,7 @@ func runOnRedisTLS(t *testing.T, fn func(client *Redis)) {
 			t.Error(err)
 		}
 		if client != nil {
-			client.Close()
+			_ = client.Close()
 		}
 	}()
 	fn(New(s.Addr(), WithTLS()))
@@ -1214,6 +1214,6 @@ type mockedNode struct {
 	RedisNode
 }
 
-func (n mockedNode) BLPop(ctx context.Context, timeout time.Duration, keys ...string) *red.StringSliceCmd {
+func (n mockedNode) BLPop(_ context.Context, _ time.Duration, _ ...string) *red.StringSliceCmd {
 	return red.NewStringSliceCmd(context.Background(), "foo", "bar")
 }

core/stores/redis/redislock.go

@@ -1,12 +1,14 @@
 package redis
 
 import (
+	"context"
 	"math/rand"
 	"strconv"
 	"sync/atomic"
 	"time"
 
 	red "github.com/go-redis/redis/v8"
+
 	"github.com/zeromicro/go-zero/core/logx"
 	"github.com/zeromicro/go-zero/core/stringx"
 )
@@ -51,8 +53,13 @@ func NewRedisLock(store *Redis, key string) *RedisLock {
 
 // Acquire acquires the lock.
 func (rl *RedisLock) Acquire() (bool, error) {
+	return rl.AcquireCtx(context.Background())
+}
+
+// AcquireCtx acquires the lock with the given ctx.
+func (rl *RedisLock) AcquireCtx(ctx context.Context) (bool, error) {
 	seconds := atomic.LoadUint32(&rl.seconds)
-	resp, err := rl.store.Eval(lockCommand, []string{rl.key}, []string{
+	resp, err := rl.store.EvalCtx(ctx, lockCommand, []string{rl.key}, []string{
 		rl.id, strconv.Itoa(int(seconds)*millisPerSecond + tolerance),
 	})
 	if err == red.Nil {
@@ -75,7 +82,12 @@ func (rl *RedisLock) Acquire() (bool, error) {
 
 // Release releases the lock.
 func (rl *RedisLock) Release() (bool, error) {
-	resp, err := rl.store.Eval(delCommand, []string{rl.key}, []string{rl.id})
+	return rl.ReleaseCtx(context.Background())
+}
+
+// ReleaseCtx releases the lock with the given ctx.
+func (rl *RedisLock) ReleaseCtx(ctx context.Context) (bool, error) {
+	resp, err := rl.store.EvalCtx(ctx, delCommand, []string{rl.key}, []string{rl.id})
 	if err != nil {
 		return false, err
 	}

core/stores/redis/redislock_test.go

@@ -1,33 +1,65 @@
 package redis
 
 import (
+	"context"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+
 	"github.com/zeromicro/go-zero/core/stringx"
 )
 
 func TestRedisLock(t *testing.T) {
-	runOnRedis(t, func(client *Redis) {
-		key := stringx.Rand()
-		firstLock := NewRedisLock(client, key)
-		firstLock.SetExpire(5)
-		firstAcquire, err := firstLock.Acquire()
-		assert.Nil(t, err)
-		assert.True(t, firstAcquire)
+	testFn := func(ctx context.Context) func(client *Redis) {
+		return func(client *Redis) {
+			key := stringx.Rand()
+			firstLock := NewRedisLock(client, key)
+			firstLock.SetExpire(5)
+			firstAcquire, err := firstLock.Acquire()
+			assert.Nil(t, err)
+			assert.True(t, firstAcquire)
 
-		secondLock := NewRedisLock(client, key)
-		secondLock.SetExpire(5)
-		againAcquire, err := secondLock.Acquire()
-		assert.Nil(t, err)
-		assert.False(t, againAcquire)
+			secondLock := NewRedisLock(client, key)
+			secondLock.SetExpire(5)
+			againAcquire, err := secondLock.Acquire()
+			assert.Nil(t, err)
+			assert.False(t, againAcquire)
 
-		release, err := firstLock.Release()
-		assert.Nil(t, err)
-		assert.True(t, release)
+			release, err := firstLock.Release()
+			assert.Nil(t, err)
+			assert.True(t, release)
 
-		endAcquire, err := secondLock.Acquire()
-		assert.Nil(t, err)
-		assert.True(t, endAcquire)
+			endAcquire, err := secondLock.Acquire()
+			assert.Nil(t, err)
+			assert.True(t, endAcquire)
+		}
+	}
+
+	t.Run("normal", func(t *testing.T) {
+		runOnRedis(t, testFn(nil))
+	})
+
+	t.Run("withContext", func(t *testing.T) {
+		runOnRedis(t, testFn(context.Background()))
+	})
+}
+
+func TestRedisLock_Expired(t *testing.T) {
+	runOnRedis(t, func(client *Redis) {
+		key := stringx.Rand()
+		redisLock := NewRedisLock(client, key)
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+		_, err := redisLock.AcquireCtx(ctx)
+		assert.NotNil(t, err)
+	})
+
+	runOnRedis(t, func(client *Redis) {
+		key := stringx.Rand()
+		redisLock := NewRedisLock(client, key)
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+		_, err := redisLock.ReleaseCtx(ctx)
+		assert.NotNil(t, err)
 	})
 }

core/stores/sqlc/cachedsql_test.go

@@ -6,7 +6,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"io/ioutil"
+	"io"
 	"log"
 	"os"
 	"runtime"
@@ -284,7 +284,7 @@ func TestCachedConn_QueryRowIndex_HasWrongCache(t *testing.T) {
 
 func TestStatCacheFails(t *testing.T) {
 	resetStats()
-	log.SetOutput(ioutil.Discard)
+	log.SetOutput(io.Discard)
 	defer log.SetOutput(os.Stdout)
 
 	r := redis.New("localhost:59999")

core/stores/sqlx/metrics.go

@@ -0,0 +1,23 @@
+package sqlx
+
+import "github.com/zeromicro/go-zero/core/metric"
+
+const namespace = "sql_client"
+
+var (
+	metricReqDur = metric.NewHistogramVec(&metric.HistogramVecOpts{
+		Namespace: namespace,
+		Subsystem: "requests",
+		Name:      "durations_ms",
+		Help:      "mysql client requests duration(ms).",
+		Labels:    []string{"command"},
+		Buckets:   []float64{5, 10, 25, 50, 100, 250, 500, 1000, 2500},
+	})
+	metricReqErr = metric.NewCounterVec(&metric.CounterVecOpts{
+		Namespace: namespace,
+		Subsystem: "requests",
+		Name:      "error_total",
+		Help:      "mysql client requests error count.",
+		Labels:    []string{"command", "error"},
+	})
+)

core/stores/sqlx/sqlconn.go

@@ -153,6 +153,9 @@ func (db *commonSqlConn) ExecCtx(ctx context.Context, q string, args ...interfac
 		result, err = exec(ctx, conn, q, args...)
 		return err
 	}, db.acceptable)
+	if err == breaker.ErrServiceUnavailable {
+		metricReqErr.Inc("Exec", "breaker")
+	}
 
 	return
 }
@@ -186,6 +189,9 @@ func (db *commonSqlConn) PrepareCtx(ctx context.Context, query string) (stmt Stm
 		}
 		return nil
 	}, db.acceptable)
+	if err == breaker.ErrServiceUnavailable {
+		metricReqErr.Inc("Prepare", "breaker")
+	}
 
 	return
 }
@@ -270,9 +276,14 @@ func (db *commonSqlConn) TransactCtx(ctx context.Context, fn func(context.Contex
 		endSpan(span, err)
 	}()
 
-	return db.brk.DoWithAcceptable(func() error {
+	err = db.brk.DoWithAcceptable(func() error {
 		return transact(ctx, db, db.beginTx, fn)
 	}, db.acceptable)
+	if err == breaker.ErrServiceUnavailable {
+		metricReqErr.Inc("Transact", "breaker")
+	}
+
+	return
 }
 
 func (db *commonSqlConn) acceptable(err error) bool {
@@ -287,7 +298,7 @@ func (db *commonSqlConn) acceptable(err error) bool {
 func (db *commonSqlConn) queryRows(ctx context.Context, scanner func(*sql.Rows) error,
 	q string, args ...interface{}) (err error) {
 	var qerr error
-	return db.brk.DoWithAcceptable(func() error {
+	err = db.brk.DoWithAcceptable(func() error {
 		conn, err := db.connProv()
 		if err != nil {
 			db.onError(err)
@@ -301,6 +312,11 @@ func (db *commonSqlConn) queryRows(ctx context.Context, scanner func(*sql.Rows)
 	}, func(err error) bool {
 		return qerr == err || db.acceptable(err)
 	})
+	if err == breaker.ErrServiceUnavailable {
+		metricReqErr.Inc("queryRows", "breaker")
+	}
+
+	return
 }
 
 func (s statement) Close() error {

core/stores/sqlx/sqlconn_test.go

@@ -17,7 +17,8 @@ func init() {
 }
 
 func TestSqlConn(t *testing.T) {
-	mock := buildConn()
+	mock, err := buildConn()
+	assert.Nil(t, err)
 	mock.ExpectExec("any")
 	mock.ExpectQuery("any").WillReturnRows(sqlmock.NewRows([]string{"foo"}))
 	conn := NewMysql(mockedDatasource)
@@ -50,8 +51,8 @@ func TestSqlConn(t *testing.T) {
 	}))
 }
 
-func buildConn() (mock sqlmock.Sqlmock) {
-	connManager.GetResource(mockedDatasource, func() (io.Closer, error) {
+func buildConn() (mock sqlmock.Sqlmock, err error) {
+	_, err = connManager.GetResource(mockedDatasource, func() (io.Closer, error) {
 		var db *sql.DB
 		var err error
 		db, mock, err = sqlmock.New()

core/stores/sqlx/stmt.go

@@ -12,7 +12,22 @@ import (
 
 const defaultSlowThreshold = time.Millisecond * 500
 
-var slowThreshold = syncx.ForAtomicDuration(defaultSlowThreshold)
+var (
+	slowThreshold = syncx.ForAtomicDuration(defaultSlowThreshold)
+	logSql        = syncx.ForAtomicBool(true)
+	logSlowSql    = syncx.ForAtomicBool(true)
+)
+
+// DisableLog disables logging of sql statements, includes info and slow logs.
+func DisableLog() {
+	logSql.Set(false)
+	logSlowSql.Set(false)
+}
+
+// DisableStmtLog disables info logging of sql statements, but keeps slow logs.
+func DisableStmtLog() {
+	logSql.Set(false)
+}
 
 // SetSlowThreshold sets the slow threshold.
 func SetSlowThreshold(threshold time.Duration) {
@@ -20,64 +35,39 @@ func SetSlowThreshold(threshold time.Duration) {
 }
 
 func exec(ctx context.Context, conn sessionConn, q string, args ...interface{}) (sql.Result, error) {
-	stmt, err := format(q, args...)
-	if err != nil {
+	guard := newGuard("exec")
+	if err := guard.start(q, args...); err != nil {
 		return nil, err
 	}
 
-	startTime := timex.Now()
 	result, err := conn.ExecContext(ctx, q, args...)
-	duration := timex.Since(startTime)
-	if duration > slowThreshold.Load() {
-		logx.WithContext(ctx).WithDuration(duration).Slowf("[SQL] exec: slowcall - %s", stmt)
-	} else {
-		logx.WithContext(ctx).WithDuration(duration).Infof("sql exec: %s", stmt)
-	}
-	if err != nil {
-		logSqlError(ctx, stmt, err)
-	}
+	guard.finish(ctx, err)
 
 	return result, err
 }
 
 func execStmt(ctx context.Context, conn stmtConn, q string, args ...interface{}) (sql.Result, error) {
-	stmt, err := format(q, args...)
-	if err != nil {
+	guard := newGuard("execStmt")
+	if err := guard.start(q, args...); err != nil {
 		return nil, err
 	}
 
-	startTime := timex.Now()
 	result, err := conn.ExecContext(ctx, args...)
-	duration := timex.Since(startTime)
-	if duration > slowThreshold.Load() {
-		logx.WithContext(ctx).WithDuration(duration).Slowf("[SQL] execStmt: slowcall - %s", stmt)
-	} else {
-		logx.WithContext(ctx).WithDuration(duration).Infof("sql execStmt: %s", stmt)
-	}
-	if err != nil {
-		logSqlError(ctx, stmt, err)
-	}
+	guard.finish(ctx, err)
 
 	return result, err
 }
 
 func query(ctx context.Context, conn sessionConn, scanner func(*sql.Rows) error,
 	q string, args ...interface{}) error {
-	stmt, err := format(q, args...)
-	if err != nil {
+	guard := newGuard("query")
+	if err := guard.start(q, args...); err != nil {
 		return err
 	}
 
-	startTime := timex.Now()
 	rows, err := conn.QueryContext(ctx, q, args...)
-	duration := timex.Since(startTime)
-	if duration > slowThreshold.Load() {
-		logx.WithContext(ctx).WithDuration(duration).Slowf("[SQL] query: slowcall - %s", stmt)
-	} else {
-		logx.WithContext(ctx).WithDuration(duration).Infof("sql query: %s", stmt)
-	}
+	guard.finish(ctx, err)
 	if err != nil {
-		logSqlError(ctx, stmt, err)
 		return err
 	}
 	defer rows.Close()
@@ -87,24 +77,76 @@ func query(ctx context.Context, conn sessionConn, scanner func(*sql.Rows) error,
 
 func queryStmt(ctx context.Context, conn stmtConn, scanner func(*sql.Rows) error,
 	q string, args ...interface{}) error {
-	stmt, err := format(q, args...)
-	if err != nil {
+	guard := newGuard("queryStmt")
+	if err := guard.start(q, args...); err != nil {
 		return err
 	}
 
-	startTime := timex.Now()
 	rows, err := conn.QueryContext(ctx, args...)
-	duration := timex.Since(startTime)
-	if duration > slowThreshold.Load() {
-		logx.WithContext(ctx).WithDuration(duration).Slowf("[SQL] queryStmt: slowcall - %s", stmt)
-	} else {
-		logx.WithContext(ctx).WithDuration(duration).Infof("sql queryStmt: %s", stmt)
-	}
+	guard.finish(ctx, err)
 	if err != nil {
-		logSqlError(ctx, stmt, err)
 		return err
 	}
 	defer rows.Close()
 
 	return scanner(rows)
 }
+
+type (
+	sqlGuard interface {
+		start(q string, args ...interface{}) error
+		finish(ctx context.Context, err error)
+	}
+
+	nilGuard struct{}
+
+	realSqlGuard struct {
+		command   string
+		stmt      string
+		startTime time.Duration
+	}
+)
+
+func newGuard(command string) sqlGuard {
+	if logSql.True() || logSlowSql.True() {
+		return &realSqlGuard{
+			command: command,
+		}
+	}
+
+	return nilGuard{}
+}
+
+func (n nilGuard) start(_ string, _ ...interface{}) error {
+	return nil
+}
+
+func (n nilGuard) finish(_ context.Context, _ error) {
+}
+
+func (e *realSqlGuard) finish(ctx context.Context, err error) {
+	duration := timex.Since(e.startTime)
+	if duration > slowThreshold.Load() {
+		logx.WithContext(ctx).WithDuration(duration).Slowf("[SQL] %s: slowcall - %s", e.command, e.stmt)
+	} else if logSql.True() {
+		logx.WithContext(ctx).WithDuration(duration).Infof("sql %s: %s", e.command, e.stmt)
+	}
+
+	if err != nil {
+		logSqlError(ctx, e.stmt, err)
+	}
+
+	metricReqDur.Observe(int64(duration/time.Millisecond), e.command)
+}
+
+func (e *realSqlGuard) start(q string, args ...interface{}) error {
+	stmt, err := format(q, args...)
+	if err != nil {
+		return err
+	}
+
+	e.stmt = stmt
+	e.startTime = timex.Now()
+
+	return nil
+}

core/stores/sqlx/stmt_test.go

@@ -178,6 +178,47 @@ func TestSetSlowThreshold(t *testing.T) {
 	assert.Equal(t, time.Second, slowThreshold.Load())
 }
 
+func TestDisableLog(t *testing.T) {
+	assert.True(t, logSql.True())
+	assert.True(t, logSlowSql.True())
+	defer func() {
+		logSql.Set(true)
+		logSlowSql.Set(true)
+	}()
+
+	DisableLog()
+	assert.False(t, logSql.True())
+	assert.False(t, logSlowSql.True())
+}
+
+func TestDisableStmtLog(t *testing.T) {
+	assert.True(t, logSql.True())
+	assert.True(t, logSlowSql.True())
+	defer func() {
+		logSql.Set(true)
+		logSlowSql.Set(true)
+	}()
+
+	DisableStmtLog()
+	assert.False(t, logSql.True())
+	assert.True(t, logSlowSql.True())
+}
+
+func TestNilGuard(t *testing.T) {
+	assert.True(t, logSql.True())
+	assert.True(t, logSlowSql.True())
+	defer func() {
+		logSql.Set(true)
+		logSlowSql.Set(true)
+	}()
+
+	DisableLog()
+	guard := newGuard("any")
+	assert.Nil(t, guard.start("foo", "bar"))
+	guard.finish(context.Background(), nil)
+	assert.Equal(t, nilGuard{}, guard)
+}
+
 type mockedSessionConn struct {
 	lastInsertId int64
 	rowsAffected int64

core/stores/sqlx/utils.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/zeromicro/go-zero/core/logx"
 	"github.com/zeromicro/go-zero/core/mapping"
@@ -152,6 +153,14 @@ func writeValue(buf *strings.Builder, arg interface{}) {
 		buf.WriteByte('\'')
 		buf.WriteString(escape(v))
 		buf.WriteByte('\'')
+	case time.Time:
+		buf.WriteByte('\'')
+		buf.WriteString(v.String())
+		buf.WriteByte('\'')
+	case *time.Time:
+		buf.WriteByte('\'')
+		buf.WriteString(v.String())
+		buf.WriteByte('\'')
 	default:
 		buf.WriteString(mapping.Repr(v))
 	}

core/stores/sqlx/utils_test.go

@@ -3,6 +3,7 @@ package sqlx
 import (
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 )
@@ -138,3 +139,14 @@ func TestFormat(t *testing.T) {
 		})
 	}
 }
+
+func TestWriteValue(t *testing.T) {
+	var buf strings.Builder
+	tm := time.Now()
+	writeValue(&buf, &tm)
+	assert.Equal(t, "'"+tm.String()+"'", buf.String())
+
+	buf.Reset()
+	writeValue(&buf, tm)
+	assert.Equal(t, "'"+tm.String()+"'", buf.String())
+}

core/syncx/resourcemanager.go

@@ -57,8 +57,8 @@ func (manager *ResourceManager) GetResource(key string, create func() (io.Closer
 		}
 
 		manager.lock.Lock()
+		defer manager.lock.Unlock()
 		manager.resources[key] = resource
-		manager.lock.Unlock()
 
 		return resource, nil
 	})

core/syncx/resourcemanager_test.go

@@ -74,6 +74,12 @@ func TestResourceManager_UseAfterClose(t *testing.T) {
 			return nil, errors.New("fail")
 		})
 		assert.NotNil(t, err)
+
+		assert.Panics(t, func() {
+			_, err = manager.GetResource("key", func() (io.Closer, error) {
+				return &dummyResource{age: 123}, nil
+			})
+		})
 	}
 }
 

core/threading/routinegroup_test.go

@@ -1,7 +1,7 @@
 package threading
 
 import (
-	"io/ioutil"
+	"io"
 	"log"
 	"sync"
 	"sync/atomic"
@@ -25,7 +25,7 @@ func TestRoutineGroupRun(t *testing.T) {
 }
 
 func TestRoutingGroupRunSafe(t *testing.T) {
-	log.SetOutput(ioutil.Discard)
+	log.SetOutput(io.Discard)
 
 	var count int32
 	group := NewRoutineGroup()

core/threading/routines_test.go

@@ -1,7 +1,7 @@
 package threading
 
 import (
-	"io/ioutil"
+	"io"
 	"log"
 	"testing"
 
@@ -14,7 +14,7 @@ func TestRoutineId(t *testing.T) {
 }
 
 func TestRunSafe(t *testing.T) {
-	log.SetOutput(ioutil.Discard)
+	log.SetOutput(io.Discard)
 
 	i := 0
 

core/trace/agent.go

@@ -1,6 +1,7 @@
 package trace
 
 import (
+	"context"
 	"fmt"
 	"sync"
 
@@ -8,21 +9,24 @@ import (
 	"github.com/zeromicro/go-zero/core/logx"
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/exporters/jaeger"
+	"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc"
 	"go.opentelemetry.io/otel/exporters/zipkin"
-	"go.opentelemetry.io/otel/propagation"
 	"go.opentelemetry.io/otel/sdk/resource"
 	sdktrace "go.opentelemetry.io/otel/sdk/trace"
 	semconv "go.opentelemetry.io/otel/semconv/v1.4.0"
+	"google.golang.org/grpc"
 )
 
 const (
 	kindJaeger = "jaeger"
 	kindZipkin = "zipkin"
+	kindGrpc   = "grpc"
 )
 
 var (
 	agents = make(map[string]lang.PlaceholderType)
 	lock   sync.Mutex
+	tp     *sdktrace.TracerProvider
 )
 
 // StartAgent starts a opentelemetry agent.
@@ -43,6 +47,11 @@ func StartAgent(c Config) {
 	agents[c.Endpoint] = lang.Placeholder
 }
 
+// StopAgent shuts down the span processors in the order they were registered.
+func StopAgent() {
+	_ = tp.Shutdown(context.Background())
+}
+
 func createExporter(c Config) (sdktrace.SpanExporter, error) {
 	// Just support jaeger and zipkin now, more for later
 	switch c.Batcher {
@@ -50,6 +59,12 @@ func createExporter(c Config) (sdktrace.SpanExporter, error) {
 		return jaeger.New(jaeger.WithCollectorEndpoint(jaeger.WithEndpoint(c.Endpoint)))
 	case kindZipkin:
 		return zipkin.New(c.Endpoint)
+	case kindGrpc:
+		return otlptracegrpc.NewUnstarted(
+			otlptracegrpc.WithInsecure(),
+			otlptracegrpc.WithEndpoint(c.Endpoint),
+			otlptracegrpc.WithDialOption(grpc.WithBlock()),
+		), nil
 	default:
 		return nil, fmt.Errorf("unknown exporter: %s", c.Batcher)
 	}
@@ -59,7 +74,7 @@ func startAgent(c Config) error {
 	opts := []sdktrace.TracerProviderOption{
 		// Set the sampling rate based on the parent span to 100%
 		sdktrace.WithSampler(sdktrace.ParentBased(sdktrace.TraceIDRatioBased(c.Sampler))),
-		// Record information about this application in an Resource.
+		// Record information about this application in a Resource.
 		sdktrace.WithResource(resource.NewSchemaless(semconv.ServiceNameKey.String(c.Name))),
 	}
 
@@ -74,10 +89,8 @@ func startAgent(c Config) error {
 		opts = append(opts, sdktrace.WithBatcher(exp))
 	}
 
-	tp := sdktrace.NewTracerProvider(opts...)
+	tp = sdktrace.NewTracerProvider(opts...)
 	otel.SetTracerProvider(tp)
-	otel.SetTextMapPropagator(propagation.NewCompositeTextMapPropagator(
-		propagation.TraceContext{}, propagation.Baggage{}))
 	otel.SetErrorHandler(otel.ErrorHandlerFunc(func(err error) {
 		logx.Errorf("[otel] error: %v", err)
 	}))

core/trace/agent_test.go

@@ -33,18 +33,24 @@ func TestStartAgent(t *testing.T) {
 		Endpoint: endpoint3,
 		Batcher:  "otlp",
 	}
+	c5 := Config{
+		Name:     "grpc",
+		Endpoint: endpoint3,
+		Batcher:  "grpc",
+	}
 
 	StartAgent(c1)
 	StartAgent(c1)
 	StartAgent(c2)
 	StartAgent(c3)
 	StartAgent(c4)
+	StartAgent(c5)
 
 	lock.Lock()
 	defer lock.Unlock()
 
 	// because remotehost cannot be resolved
-	assert.Equal(t, 2, len(agents))
+	assert.Equal(t, 3, len(agents))
 	_, ok := agents[""]
 	assert.True(t, ok)
 	_, ok = agents[endpoint1]

core/trace/config.go

@@ -8,5 +8,5 @@ type Config struct {
 	Name     string  `json:",optional"`
 	Endpoint string  `json:",optional"`
 	Sampler  float64 `json:",default=1.0"`
-	Batcher  string  `json:",default=jaeger,options=jaeger|zipkin"`
+	Batcher  string  `json:",default=jaeger,options=jaeger|zipkin|grpc"`
 }

core/trace/propagation.go

@@ -0,0 +1,11 @@
+package trace
+
+import (
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/propagation"
+)
+
+func init() {
+	otel.SetTextMapPropagator(propagation.NewCompositeTextMapPropagator(
+		propagation.TraceContext{}, propagation.Baggage{}))
+}

core/trace/tracer.go

@@ -38,7 +38,7 @@ func (s *metadataSupplier) Keys() []string {
 	return out
 }
 
-// Inject injects the metadata into ctx.
+// Inject injects cross-cutting concerns from the ctx into the metadata.
 func Inject(ctx context.Context, p propagation.TextMapPropagator, metadata *metadata.MD) {
 	p.Inject(ctx, &metadataSupplier{
 		metadata: metadata,

core/utils/version.go

@@ -35,13 +35,11 @@ func CompareVersions(v1, op, v2 string) bool {
 
 // return -1 if v1<v2, 0 if they are equal, and 1 if v1>v2
 func compare(v1, v2 string) int {
-	v1 = replacer.Replace(v1)
-	v2 = replacer.Replace(v2)
-	fields1 := strings.Split(v1, ".")
-	fields2 := strings.Split(v2, ".")
-	ver1 := strsToInts(fields1)
-	ver2 := strsToInts(fields2)
-	shorter := mathx.MinInt(len(ver1), len(ver2))
+	v1, v2 = replacer.Replace(v1), replacer.Replace(v2)
+	fields1, fields2 := strings.Split(v1, "."), strings.Split(v2, ".")
+	ver1, ver2 := strsToInts(fields1), strsToInts(fields2)
+	ver1len, ver2len := len(ver1), len(ver2)
+	shorter := mathx.MinInt(ver1len, ver2len)
 
 	for i := 0; i < shorter; i++ {
 		if ver1[i] == ver2[i] {
@@ -53,9 +51,9 @@ func compare(v1, v2 string) int {
 		}
 	}
 
-	if len(ver1) < len(ver2) {
+	if ver1len < ver2len {
 		return -1
-	} else if len(ver1) == len(ver2) {
+	} else if ver1len == ver2len {
 		return 0
 	} else {
 		return 1

gateway/config.go

@@ -0,0 +1,42 @@
+package gateway
+
+import (
+	"time"
+
+	"github.com/zeromicro/go-zero/rest"
+	"github.com/zeromicro/go-zero/zrpc"
+)
+
+type (
+	// GatewayConf is the configuration for gateway.
+	GatewayConf struct {
+		rest.RestConf
+		Upstreams []Upstream
+		Timeout   time.Duration `json:",default=5s"`
+	}
+
+	// RouteMapping is a mapping between a gateway route and an upstream rpc method.
+	RouteMapping struct {
+		// Method is the HTTP method, like GET, POST, PUT, DELETE.
+		Method string
+		// Path is the HTTP path.
+		Path string
+		// RpcPath is the gRPC rpc method, with format of package.service/method
+		RpcPath string
+	}
+
+	// Upstream is the configuration for an upstream.
+	Upstream struct {
+		// Name is the name of the upstream.
+		Name string `json:",optional"`
+		// Grpc is the target of the upstream.
+		Grpc zrpc.RpcClientConf
+		// ProtoSets is the file list of proto set, like [hello.pb].
+		// if your proto file import another proto file, you need to write multi-file slice,
+		// like [hello.pb, common.pb].
+		ProtoSets []string `json:",optional"`
+		// Mappings is the mapping between gateway routes and Upstream rpc methods.
+		// Keep it blank if annotations are added in rpc methods.
+		Mappings []RouteMapping `json:",optional"`
+	}
+)

gateway/internal/descriptorsource.go

@@ -0,0 +1,102 @@
+package internal
+
+import (
+	"fmt"
+	"net/http"
+	"strings"
+
+	"github.com/fullstorydev/grpcurl"
+	"github.com/jhump/protoreflect/desc"
+	"google.golang.org/genproto/googleapis/api/annotations"
+	"google.golang.org/protobuf/proto"
+)
+
+type Method struct {
+	HttpMethod string
+	HttpPath   string
+	RpcPath    string
+}
+
+// GetMethods returns all methods of the given grpcurl.DescriptorSource.
+func GetMethods(source grpcurl.DescriptorSource) ([]Method, error) {
+	svcs, err := source.ListServices()
+	if err != nil {
+		return nil, err
+	}
+
+	var methods []Method
+	for _, svc := range svcs {
+		d, err := source.FindSymbol(svc)
+		if err != nil {
+			return nil, err
+		}
+
+		switch val := d.(type) {
+		case *desc.ServiceDescriptor:
+			svcMethods := val.GetMethods()
+			for _, method := range svcMethods {
+				rpcPath := fmt.Sprintf("%s/%s", svc, method.GetName())
+				ext := proto.GetExtension(method.GetMethodOptions(), annotations.E_Http)
+				if ext == nil {
+					methods = append(methods, Method{
+						RpcPath: rpcPath,
+					})
+					continue
+				}
+
+				httpExt, ok := ext.(*annotations.HttpRule)
+				if !ok {
+					methods = append(methods, Method{
+						RpcPath: rpcPath,
+					})
+					continue
+				}
+
+				switch rule := httpExt.GetPattern().(type) {
+				case *annotations.HttpRule_Get:
+					methods = append(methods, Method{
+						HttpMethod: http.MethodGet,
+						HttpPath:   adjustHttpPath(rule.Get),
+						RpcPath:    rpcPath,
+					})
+				case *annotations.HttpRule_Post:
+					methods = append(methods, Method{
+						HttpMethod: http.MethodPost,
+						HttpPath:   adjustHttpPath(rule.Post),
+						RpcPath:    rpcPath,
+					})
+				case *annotations.HttpRule_Put:
+					methods = append(methods, Method{
+						HttpMethod: http.MethodPut,
+						HttpPath:   adjustHttpPath(rule.Put),
+						RpcPath:    rpcPath,
+					})
+				case *annotations.HttpRule_Delete:
+					methods = append(methods, Method{
+						HttpMethod: http.MethodDelete,
+						HttpPath:   adjustHttpPath(rule.Delete),
+						RpcPath:    rpcPath,
+					})
+				case *annotations.HttpRule_Patch:
+					methods = append(methods, Method{
+						HttpMethod: http.MethodPatch,
+						HttpPath:   adjustHttpPath(rule.Patch),
+						RpcPath:    rpcPath,
+					})
+				default:
+					methods = append(methods, Method{
+						RpcPath: rpcPath,
+					})
+				}
+			}
+		}
+	}
+
+	return methods, nil
+}
+
+func adjustHttpPath(path string) string {
+	path = strings.ReplaceAll(path, "{", ":")
+	path = strings.ReplaceAll(path, "}", "")
+	return path
+}

gateway/internal/eventhandler.go

@@ -0,0 +1,47 @@
+package internal
+
+import (
+	"io"
+
+	"github.com/golang/protobuf/jsonpb"
+	"github.com/golang/protobuf/proto"
+	"github.com/jhump/protoreflect/desc"
+	"github.com/zeromicro/go-zero/core/logx"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/status"
+)
+
+type EventHandler struct {
+	Status    *status.Status
+	writer    io.Writer
+	marshaler jsonpb.Marshaler
+}
+
+func NewEventHandler(writer io.Writer, resolver jsonpb.AnyResolver) *EventHandler {
+	return &EventHandler{
+		writer: writer,
+		marshaler: jsonpb.Marshaler{
+			EmitDefaults: true,
+			AnyResolver:  resolver,
+		},
+	}
+}
+
+func (h *EventHandler) OnResolveMethod(_ *desc.MethodDescriptor) {
+}
+
+func (h *EventHandler) OnSendHeaders(_ metadata.MD) {
+}
+
+func (h *EventHandler) OnReceiveHeaders(_ metadata.MD) {
+}
+
+func (h *EventHandler) OnReceiveResponse(message proto.Message) {
+	if err := h.marshaler.Marshal(h.writer, message); err != nil {
+		logx.Error(err)
+	}
+}
+
+func (h *EventHandler) OnReceiveTrailers(status *status.Status, _ metadata.MD) {
+	h.Status = status
+}

gateway/internal/headerprocessor.go

@@ -0,0 +1,30 @@
+package internal
+
+import (
+	"fmt"
+	"net/http"
+	"strings"
+)
+
+const (
+	metadataHeaderPrefix = "Grpc-Metadata-"
+	metadataPrefix       = "gateway-"
+)
+
+// ProcessHeaders builds the headers for the gateway from HTTP headers.
+func ProcessHeaders(header http.Header) []string {
+	var headers []string
+
+	for k, v := range header {
+		if !strings.HasPrefix(k, metadataHeaderPrefix) {
+			continue
+		}
+
+		key := fmt.Sprintf("%s%s", metadataPrefix, strings.TrimPrefix(k, metadataHeaderPrefix))
+		for _, vv := range v {
+			headers = append(headers, key+":"+vv)
+		}
+	}
+
+	return headers
+}

gateway/internal/headerprocessor_test.go

@@ -0,0 +1,21 @@
+package internal
+
+import (
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestBuildHeadersNoValue(t *testing.T) {
+	req := httptest.NewRequest("GET", "/", nil)
+	req.Header.Add("a", "b")
+	assert.Nil(t, ProcessHeaders(req.Header))
+}
+
+func TestBuildHeadersWithValues(t *testing.T) {
+	req := httptest.NewRequest("GET", "/", nil)
+	req.Header.Add("grpc-metadata-a", "b")
+	req.Header.Add("grpc-metadata-b", "b")
+	assert.ElementsMatch(t, []string{"gateway-A:b", "gateway-B:b"}, ProcessHeaders(req.Header))
+}

gateway/internal/requestparser.go

@@ -0,0 +1,53 @@
+package internal
+
+import (
+	"bytes"
+	"encoding/json"
+	"net/http"
+
+	"github.com/fullstorydev/grpcurl"
+	"github.com/golang/protobuf/jsonpb"
+	"github.com/zeromicro/go-zero/rest/httpx"
+	"github.com/zeromicro/go-zero/rest/pathvar"
+)
+
+// NewRequestParser creates a new request parser from the given http.Request and resolver.
+func NewRequestParser(r *http.Request, resolver jsonpb.AnyResolver) (grpcurl.RequestParser, error) {
+	vars := pathvar.Vars(r)
+	params, err := httpx.GetFormValues(r)
+	if err != nil {
+		return nil, err
+	}
+
+	for k, v := range vars {
+		params[k] = v
+	}
+	if len(params) == 0 {
+		return grpcurl.NewJSONRequestParser(r.Body, resolver), nil
+	}
+
+	if r.ContentLength == 0 {
+		return buildJsonRequestParser(params, resolver)
+	}
+
+	m := make(map[string]interface{})
+	if err := json.NewDecoder(r.Body).Decode(&m); err != nil {
+		return nil, err
+	}
+
+	for k, v := range params {
+		m[k] = v
+	}
+
+	return buildJsonRequestParser(m, resolver)
+}
+
+func buildJsonRequestParser(m map[string]interface{}, resolver jsonpb.AnyResolver) (
+	grpcurl.RequestParser, error) {
+	var buf bytes.Buffer
+	if err := json.NewEncoder(&buf).Encode(m); err != nil {
+		return nil, err
+	}
+
+	return grpcurl.NewJSONRequestParser(&buf, resolver), nil
+}

gateway/internal/requestparser_test.go

@@ -0,0 +1,55 @@
+package internal
+
+import (
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/zeromicro/go-zero/rest/pathvar"
+)
+
+func TestNewRequestParserNoVar(t *testing.T) {
+	req := httptest.NewRequest("GET", "/", nil)
+	parser, err := NewRequestParser(req, nil)
+	assert.Nil(t, err)
+	assert.NotNil(t, parser)
+}
+
+func TestNewRequestParserWithVars(t *testing.T) {
+	req := httptest.NewRequest("GET", "/", nil)
+	req = pathvar.WithVars(req, map[string]string{"a": "b"})
+	parser, err := NewRequestParser(req, nil)
+	assert.Nil(t, err)
+	assert.NotNil(t, parser)
+}
+
+func TestNewRequestParserNoVarWithBody(t *testing.T) {
+	req := httptest.NewRequest("GET", "/", strings.NewReader(`{"a": "b"}`))
+	parser, err := NewRequestParser(req, nil)
+	assert.Nil(t, err)
+	assert.NotNil(t, parser)
+}
+
+func TestNewRequestParserWithVarsWithBody(t *testing.T) {
+	req := httptest.NewRequest("GET", "/", strings.NewReader(`{"a": "b"}`))
+	req = pathvar.WithVars(req, map[string]string{"c": "d"})
+	parser, err := NewRequestParser(req, nil)
+	assert.Nil(t, err)
+	assert.NotNil(t, parser)
+}
+
+func TestNewRequestParserWithVarsWithWrongBody(t *testing.T) {
+	req := httptest.NewRequest("GET", "/", strings.NewReader(`{"a": "b"`))
+	req = pathvar.WithVars(req, map[string]string{"c": "d"})
+	parser, err := NewRequestParser(req, nil)
+	assert.NotNil(t, err)
+	assert.Nil(t, parser)
+}
+
+func TestNewRequestParserWithForm(t *testing.T) {
+	req := httptest.NewRequest("GET", "/val?a=b", nil)
+	parser, err := NewRequestParser(req, nil)
+	assert.Nil(t, err)
+	assert.NotNil(t, parser)
+}

gateway/internal/timeout.go

@@ -0,0 +1,19 @@
+package internal
+
+import (
+	"net/http"
+	"time"
+)
+
+const grpcTimeoutHeader = "Grpc-Timeout"
+
+// GetTimeout returns the timeout from the header, if not set, returns the default timeout.
+func GetTimeout(header http.Header, defaultTimeout time.Duration) time.Duration {
+	if timeout := header.Get(grpcTimeoutHeader); len(timeout) > 0 {
+		if t, err := time.ParseDuration(timeout); err == nil {
+			return t
+		}
+	}
+
+	return defaultTimeout
+}

gateway/internal/timeout_test.go

@@ -0,0 +1,22 @@
+package internal
+
+import (
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGetTimeout(t *testing.T) {
+	req := httptest.NewRequest("GET", "/", nil)
+	req.Header.Set(grpcTimeoutHeader, "1s")
+	timeout := GetTimeout(req.Header, time.Second*5)
+	assert.Equal(t, time.Second, timeout)
+}
+
+func TestGetTimeoutDefault(t *testing.T) {
+	req := httptest.NewRequest("GET", "/", nil)
+	timeout := GetTimeout(req.Header, time.Second*5)
+	assert.Equal(t, time.Second*5, timeout)
+}

gateway/readme.md

@@ -0,0 +1,63 @@
+# Gateway
+
+## Usage
+
+- main.go
+
+```go
+var configFile = flag.String("f", "config.yaml", "config file")
+
+func main() {
+    flag.Parse()
+
+    var c gateway.GatewayConf
+    conf.MustLoad(*configFile, &c)
+    gw := gateway.MustNewServer(c)
+    defer gw.Stop()
+    gw.Start()
+}
+```
+
+- config.yaml
+
+```yaml
+Name: demo-gateway
+Host: localhost
+Port: 8888
+Upstreams:
+  - Grpc:
+      Etcd:
+        Hosts:
+          - localhost:2379
+        Key: hello.rpc
+    # protoset mode
+    ProtoSets:
+      - hello.pb
+    # Mappings can also be written in proto options
+    Mappings:
+      - Method: get
+        Path: /pingHello/:ping
+        RpcPath: hello.Hello/Ping
+  - Grpc:
+      Endpoints:
+        - localhost:8081
+    # reflection mode, no ProtoSet settings
+    Mappings:
+      - Method: post
+        Path: /pingWorld
+        RpcPath: world.World/Ping
+```
+
+## Generate ProtoSet files
+
+- example command without external imports
+
+```shell
+protoc --descriptor_set_out=hello.pb hello.proto
+```
+
+- example command with external imports
+
+```shell
+protoc --include_imports --proto_path=. --descriptor_set_out=hello.pb hello.proto
+```

gateway/server.go

@@ -0,0 +1,193 @@
+package gateway
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"strings"
+	"time"
+
+	"github.com/fullstorydev/grpcurl"
+	"github.com/golang/protobuf/jsonpb"
+	"github.com/jhump/protoreflect/grpcreflect"
+	"github.com/zeromicro/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/core/mr"
+	"github.com/zeromicro/go-zero/gateway/internal"
+	"github.com/zeromicro/go-zero/rest"
+	"github.com/zeromicro/go-zero/rest/httpx"
+	"github.com/zeromicro/go-zero/zrpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/reflection/grpc_reflection_v1alpha"
+)
+
+type (
+	// Server is a gateway server.
+	Server struct {
+		*rest.Server
+		upstreams     []Upstream
+		timeout       time.Duration
+		processHeader func(http.Header) []string
+	}
+
+	// Option defines the method to customize Server.
+	Option func(svr *Server)
+)
+
+// MustNewServer creates a new gateway server.
+func MustNewServer(c GatewayConf, opts ...Option) *Server {
+	svr := &Server{
+		Server:    rest.MustNewServer(c.RestConf),
+		upstreams: c.Upstreams,
+		timeout:   c.Timeout,
+	}
+	for _, opt := range opts {
+		opt(svr)
+	}
+
+	return svr
+}
+
+// Start starts the gateway server.
+func (s *Server) Start() {
+	logx.Must(s.build())
+	s.Server.Start()
+}
+
+// Stop stops the gateway server.
+func (s *Server) Stop() {
+	s.Server.Stop()
+}
+
+func (s *Server) build() error {
+	if err := s.ensureUpstreamNames(); err != nil {
+		return err
+	}
+
+	return mr.MapReduceVoid(func(source chan<- interface{}) {
+		for _, up := range s.upstreams {
+			source <- up
+		}
+	}, func(item interface{}, writer mr.Writer, cancel func(error)) {
+		up := item.(Upstream)
+		cli := zrpc.MustNewClient(up.Grpc)
+		source, err := s.createDescriptorSource(cli, up)
+		if err != nil {
+			cancel(fmt.Errorf("%s: %w", up.Name, err))
+			return
+		}
+
+		methods, err := internal.GetMethods(source)
+		if err != nil {
+			cancel(fmt.Errorf("%s: %w", up.Name, err))
+			return
+		}
+
+		resolver := grpcurl.AnyResolverFromDescriptorSource(source)
+		for _, m := range methods {
+			if len(m.HttpMethod) > 0 && len(m.HttpPath) > 0 {
+				writer.Write(rest.Route{
+					Method:  m.HttpMethod,
+					Path:    m.HttpPath,
+					Handler: s.buildHandler(source, resolver, cli, m.RpcPath),
+				})
+			}
+		}
+
+		methodSet := make(map[string]struct{})
+		for _, m := range methods {
+			methodSet[m.RpcPath] = struct{}{}
+		}
+		for _, m := range up.Mappings {
+			if _, ok := methodSet[m.RpcPath]; !ok {
+				cancel(fmt.Errorf("%s: rpc method %s not found", up.Name, m.RpcPath))
+				return
+			}
+
+			writer.Write(rest.Route{
+				Method:  strings.ToUpper(m.Method),
+				Path:    m.Path,
+				Handler: s.buildHandler(source, resolver, cli, m.RpcPath),
+			})
+		}
+	}, func(pipe <-chan interface{}, cancel func(error)) {
+		for item := range pipe {
+			route := item.(rest.Route)
+			s.Server.AddRoute(route)
+		}
+	})
+}
+
+func (s *Server) buildHandler(source grpcurl.DescriptorSource, resolver jsonpb.AnyResolver,
+	cli zrpc.Client, rpcPath string) func(http.ResponseWriter, *http.Request) {
+	return func(w http.ResponseWriter, r *http.Request) {
+		parser, err := internal.NewRequestParser(r, resolver)
+		if err != nil {
+			httpx.Error(w, err)
+			return
+		}
+
+		timeout := internal.GetTimeout(r.Header, s.timeout)
+		ctx, can := context.WithTimeout(r.Context(), timeout)
+		defer can()
+
+		w.Header().Set(httpx.ContentType, httpx.JsonContentType)
+		handler := internal.NewEventHandler(w, resolver)
+		if err := grpcurl.InvokeRPC(ctx, source, cli.Conn(), rpcPath, s.prepareMetadata(r.Header),
+			handler, parser.Next); err != nil {
+			httpx.Error(w, err)
+		}
+
+		st := handler.Status
+		if st.Code() != codes.OK {
+			httpx.Error(w, st.Err())
+		}
+	}
+}
+
+func (s *Server) createDescriptorSource(cli zrpc.Client, up Upstream) (grpcurl.DescriptorSource, error) {
+	var source grpcurl.DescriptorSource
+	var err error
+
+	if len(up.ProtoSets) > 0 {
+		source, err = grpcurl.DescriptorSourceFromProtoSets(up.ProtoSets...)
+		if err != nil {
+			return nil, err
+		}
+	} else {
+		refCli := grpc_reflection_v1alpha.NewServerReflectionClient(cli.Conn())
+		client := grpcreflect.NewClient(context.Background(), refCli)
+		source = grpcurl.DescriptorSourceFromServer(context.Background(), client)
+	}
+
+	return source, nil
+}
+
+func (s *Server) ensureUpstreamNames() error {
+	for _, up := range s.upstreams {
+		target, err := up.Grpc.BuildTarget()
+		if err != nil {
+			return err
+		}
+
+		up.Name = target
+	}
+
+	return nil
+}
+
+func (s *Server) prepareMetadata(header http.Header) []string {
+	vals := internal.ProcessHeaders(header)
+	if s.processHeader != nil {
+		vals = append(vals, s.processHeader(header)...)
+	}
+
+	return vals
+}
+
+// WithHeaderProcessor sets a processor to process request headers.
+// The returned headers are used as metadata to invoke the RPC.
+func WithHeaderProcessor(processHeader func(http.Header) []string) func(*Server) {
+	return func(s *Server) {
+		s.processHeader = processHeader
+	}
+}

go.mod

@@ -5,42 +5,45 @@ go 1.16
 require (
 	github.com/ClickHouse/clickhouse-go/v2 v2.0.14
 	github.com/DATA-DOG/go-sqlmock v1.5.0
-	github.com/alicebob/miniredis/v2 v2.21.0
+	github.com/alicebob/miniredis/v2 v2.23.0
 	github.com/fatih/color v1.13.0
+	github.com/fullstorydev/grpcurl v1.8.7
 	github.com/globalsign/mgo v0.0.0-20181015135952-eeefdecb41b8
 	github.com/go-redis/redis/v8 v8.11.5
 	github.com/go-sql-driver/mysql v1.6.0
-	github.com/golang-jwt/jwt/v4 v4.4.1
+	github.com/golang-jwt/jwt/v4 v4.4.2
 	github.com/golang/mock v1.6.0
+	github.com/golang/protobuf v1.5.2
 	github.com/google/uuid v1.3.0
-	github.com/justinas/alice v1.2.0
-	github.com/lib/pq v1.10.6
+	github.com/jhump/protoreflect v1.13.0
+	github.com/lib/pq v1.10.7
 	github.com/olekukonko/tablewriter v0.0.5
-	github.com/pelletier/go-toml/v2 v2.0.1
-	github.com/prometheus/client_golang v1.12.2
+	github.com/pelletier/go-toml/v2 v2.0.5
+	github.com/prometheus/client_golang v1.13.0
 	github.com/spaolacci/murmur3 v1.1.0
-	github.com/stretchr/testify v1.7.1
-	go.etcd.io/etcd/api/v3 v3.5.4
-	go.etcd.io/etcd/client/v3 v3.5.4
-	go.mongodb.org/mongo-driver v1.9.1
-	go.opentelemetry.io/otel v1.7.0
-	go.opentelemetry.io/otel/exporters/jaeger v1.7.0
-	go.opentelemetry.io/otel/exporters/zipkin v1.7.0
-	go.opentelemetry.io/otel/sdk v1.7.0
-	go.opentelemetry.io/otel/trace v1.7.0
+	github.com/stretchr/testify v1.8.0
+	go.etcd.io/etcd/api/v3 v3.5.5
+	go.etcd.io/etcd/client/v3 v3.5.5
+	go.mongodb.org/mongo-driver v1.10.2
+	go.opentelemetry.io/otel v1.10.0
+	go.opentelemetry.io/otel/exporters/jaeger v1.10.0
+	go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.10.0
+	go.opentelemetry.io/otel/exporters/zipkin v1.10.0
+	go.opentelemetry.io/otel/sdk v1.10.0
+	go.opentelemetry.io/otel/trace v1.10.0
 	go.uber.org/automaxprocs v1.5.1
-	go.uber.org/goleak v1.1.12
+	go.uber.org/goleak v1.2.0
 	golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a
 	golang.org/x/time v0.0.0-20220411224347-583f2d630306
-	google.golang.org/grpc v1.47.0
-	google.golang.org/protobuf v1.28.0
+	google.golang.org/grpc v1.49.0
+	google.golang.org/protobuf v1.28.1
 	gopkg.in/cheggaaa/pb.v1 v1.0.28
 	gopkg.in/h2non/gock.v1 v1.1.2
 	gopkg.in/yaml.v2 v2.4.0
 	k8s.io/api v0.22.9
 	k8s.io/apimachinery v0.22.9
 	k8s.io/client-go v0.22.9
-	k8s.io/utils v0.0.0-20220210201930-3a6ce19ff2f9
+	k8s.io/utils v0.0.0-20220706174534-f6158b442e7c
 )
 
 require (
@@ -51,6 +54,6 @@ require (
 	go.uber.org/multierr v1.8.0 // indirect
 	go.uber.org/zap v1.21.0 // indirect
 	golang.org/x/net v0.0.0-20220531201128-c960675eff93 // indirect
-	google.golang.org/genproto v0.0.0-20220602131408-e326c6e8e9c8 // indirect
+	google.golang.org/genproto v0.0.0-20220602131408-e326c6e8e9c8
 	k8s.io/klog/v2 v2.40.1 // indirect
 )