refactor: move json related header vars to internal (#1840)

* refactor: move json related header vars to internal

* refactor: use header.ContentType

.github/workflows/go.yml

@@ -28,7 +28,7 @@ jobs:
         run: |
           go vet -stdmethods=false $(go list ./...)
           go install mvdan.cc/gofumpt@latest
-          test -z "$(gofumpt -s -l -extra .)" || echo "Please run 'gofumpt -l -w -extra .'"
+          test -z "$(gofumpt -l -extra .)" || echo "Please run 'gofumpt -l -w -extra .'"
 
       - name: Test
         run: go test -race -coverprofile=coverage.txt -covermode=atomic ./...

.gitignore

@@ -17,7 +17,7 @@
 
 # for test purpose
 **/adhoc
-**/testdata
+go.work
 
 # gitlab ci
 .cache

LICENSE

@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2020 xiaoheiban_server_go
+Copyright (c) 2022 zeromicro
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

core/collection/cache.go

@@ -98,13 +98,18 @@ func (c *Cache) Get(key string) (interface{}, bool) {
 
 // Set sets value into c with key.
 func (c *Cache) Set(key string, value interface{}) {
+	c.SetWithExpire(key, value, c.expire)
+}
+
+// SetWithExpire sets value into c with key and expire with the given value.
+func (c *Cache) SetWithExpire(key string, value interface{}, expire time.Duration) {
 	c.lock.Lock()
 	_, ok := c.data[key]
 	c.data[key] = value
 	c.lruCache.add(key)
 	c.lock.Unlock()
 
-	expiry := c.unstableExpiry.AroundDuration(c.expire)
+	expiry := c.unstableExpiry.AroundDuration(expire)
 	if ok {
 		c.timingWheel.MoveTimer(key, expiry)
 	} else {

core/collection/cache_test.go

@@ -18,7 +18,7 @@ func TestCacheSet(t *testing.T) {
 	assert.Nil(t, err)
 
 	cache.Set("first", "first element")
-	cache.Set("second", "second element")
+	cache.SetWithExpire("second", "second element", time.Second*3)
 
 	value, ok := cache.Get("first")
 	assert.True(t, ok)

core/collection/timingwheel.go

@@ -2,6 +2,7 @@ package collection
 
 import (
 	"container/list"
+	"errors"
 	"fmt"
 	"time"
 
@@ -12,6 +13,11 @@ import (
 
 const drainWorkers = 8
 
+var (
+	ErrClosed   = errors.New("TimingWheel is closed already")
+	ErrArgument = errors.New("incorrect task argument")
+)
+
 type (
 	// Execute defines the method to execute the task.
 	Execute func(key, value interface{})
@@ -59,14 +65,15 @@ type (
 // NewTimingWheel returns a TimingWheel.
 func NewTimingWheel(interval time.Duration, numSlots int, execute Execute) (*TimingWheel, error) {
 	if interval <= 0 || numSlots <= 0 || execute == nil {
-		return nil, fmt.Errorf("interval: %v, slots: %d, execute: %p", interval, numSlots, execute)
+		return nil, fmt.Errorf("interval: %v, slots: %d, execute: %p",
+			interval, numSlots, execute)
 	}
 
 	return newTimingWheelWithClock(interval, numSlots, execute, timex.NewTicker(interval))
 }
 
-func newTimingWheelWithClock(interval time.Duration, numSlots int, execute Execute, ticker timex.Ticker) (
-	*TimingWheel, error) {
+func newTimingWheelWithClock(interval time.Duration, numSlots int, execute Execute,
+	ticker timex.Ticker) (*TimingWheel, error) {
 	tw := &TimingWheel{
 		interval:      interval,
 		ticker:        ticker,
@@ -89,47 +96,67 @@ func newTimingWheelWithClock(interval time.Duration, numSlots int, execute Execu
 }
 
 // Drain drains all items and executes them.
-func (tw *TimingWheel) Drain(fn func(key, value interface{})) {
-	tw.drainChannel <- fn
+func (tw *TimingWheel) Drain(fn func(key, value interface{})) error {
+	select {
+	case tw.drainChannel <- fn:
+		return nil
+	case <-tw.stopChannel:
+		return ErrClosed
+	}
 }
 
 // MoveTimer moves the task with the given key to the given delay.
-func (tw *TimingWheel) MoveTimer(key interface{}, delay time.Duration) {
+func (tw *TimingWheel) MoveTimer(key interface{}, delay time.Duration) error {
 	if delay <= 0 || key == nil {
-		return
+		return ErrArgument
 	}
 
-	tw.moveChannel <- baseEntry{
+	select {
+	case tw.moveChannel <- baseEntry{
 		delay: delay,
 		key:   key,
+	}:
+		return nil
+	case <-tw.stopChannel:
+		return ErrClosed
 	}
 }
 
 // RemoveTimer removes the task with the given key.
-func (tw *TimingWheel) RemoveTimer(key interface{}) {
+func (tw *TimingWheel) RemoveTimer(key interface{}) error {
 	if key == nil {
-		return
+		return ErrArgument
 	}
 
-	tw.removeChannel <- key
+	select {
+	case tw.removeChannel <- key:
+		return nil
+	case <-tw.stopChannel:
+		return ErrClosed
+	}
 }
 
 // SetTimer sets the task value with the given key to the delay.
-func (tw *TimingWheel) SetTimer(key, value interface{}, delay time.Duration) {
+func (tw *TimingWheel) SetTimer(key, value interface{}, delay time.Duration) error {
 	if delay <= 0 || key == nil {
-		return
+		return ErrArgument
 	}
 
-	tw.setChannel <- timingEntry{
+	select {
+	case tw.setChannel <- timingEntry{
 		baseEntry: baseEntry{
 			delay: delay,
 			key:   key,
 		},
 		value: value,
+	}:
+		return nil
+	case <-tw.stopChannel:
+		return ErrClosed
 	}
 }
 
-// Stop stops tw.
+// Stop stops tw. No more actions after stopping a TimingWheel.
 func (tw *TimingWheel) Stop() {
 	close(tw.stopChannel)
 }

core/collection/timingwheel_test.go

@@ -28,7 +28,6 @@ func TestTimingWheel_Drain(t *testing.T) {
 	ticker := timex.NewFakeTicker()
 	tw, _ := newTimingWheelWithClock(testStep, 10, func(k, v interface{}) {
 	}, ticker)
-	defer tw.Stop()
 	tw.SetTimer("first", 3, testStep*4)
 	tw.SetTimer("second", 5, testStep*7)
 	tw.SetTimer("third", 7, testStep*7)
@@ -56,6 +55,8 @@ func TestTimingWheel_Drain(t *testing.T) {
 	})
 	time.Sleep(time.Millisecond * 100)
 	assert.Equal(t, 0, count)
+	tw.Stop()
+	assert.Equal(t, ErrClosed, tw.Drain(func(key, value interface{}) {}))
 }
 
 func TestTimingWheel_SetTimerSoon(t *testing.T) {
@@ -102,6 +103,13 @@ func TestTimingWheel_SetTimerWrongDelay(t *testing.T) {
 	})
 }
 
+func TestTimingWheel_SetTimerAfterClose(t *testing.T) {
+	ticker := timex.NewFakeTicker()
+	tw, _ := newTimingWheelWithClock(testStep, 10, func(k, v interface{}) {}, ticker)
+	tw.Stop()
+	assert.Equal(t, ErrClosed, tw.SetTimer("any", 3, testStep))
+}
+
 func TestTimingWheel_MoveTimer(t *testing.T) {
 	run := syncx.NewAtomicBool()
 	ticker := timex.NewFakeTicker()
@@ -111,7 +119,6 @@ func TestTimingWheel_MoveTimer(t *testing.T) {
 		assert.Equal(t, 3, v.(int))
 		ticker.Done()
 	}, ticker)
-	defer tw.Stop()
 	tw.SetTimer("any", 3, testStep*4)
 	tw.MoveTimer("any", testStep*7)
 	tw.MoveTimer("any", -testStep)
@@ -125,6 +132,8 @@ func TestTimingWheel_MoveTimer(t *testing.T) {
 	}
 	assert.Nil(t, ticker.Wait(waitTime))
 	assert.True(t, run.True())
+	tw.Stop()
+	assert.Equal(t, ErrClosed, tw.MoveTimer("any", time.Millisecond))
 }
 
 func TestTimingWheel_MoveTimerSoon(t *testing.T) {
@@ -175,6 +184,7 @@ func TestTimingWheel_RemoveTimer(t *testing.T) {
 		ticker.Tick()
 	}
 	tw.Stop()
+	assert.Equal(t, ErrClosed, tw.RemoveTimer("any"))
 }
 
 func TestTimingWheel_SetTimer(t *testing.T) {

core/logx/logs.go

@@ -275,7 +275,7 @@ func Infov(v interface{}) {
 	infoAnySync(v)
 }
 
-// Must checks if err is nil, otherwise logs the err and exits.
+// Must checks if err is nil, otherwise logs the error and exits.
 func Must(err error) {
 	if err != nil {
 		msg := formatWithCaller(err.Error(), 3)

core/mapping/marshaler.go

@@ -0,0 +1,186 @@
+package mapping
+
+import (
+	"fmt"
+	"reflect"
+	"strings"
+)
+
+const (
+	emptyTag       = ""
+	tagKVSeparator = ":"
+)
+
+// Marshal marshals the given val and returns the map that contains the fields.
+// optional=another is not implemented, and it's hard to implement and not common used.
+func Marshal(val interface{}) (map[string]map[string]interface{}, error) {
+	ret := make(map[string]map[string]interface{})
+	tp := reflect.TypeOf(val)
+	if tp.Kind() == reflect.Ptr {
+		tp = tp.Elem()
+	}
+	rv := reflect.ValueOf(val)
+	if rv.Kind() == reflect.Ptr {
+		rv = rv.Elem()
+	}
+
+	for i := 0; i < tp.NumField(); i++ {
+		field := tp.Field(i)
+		value := rv.Field(i)
+		if err := processMember(field, value, ret); err != nil {
+			return nil, err
+		}
+	}
+
+	return ret, nil
+}
+
+func getTag(field reflect.StructField) (string, bool) {
+	tag := string(field.Tag)
+	if i := strings.Index(tag, tagKVSeparator); i >= 0 {
+		return strings.TrimSpace(tag[:i]), true
+	}
+
+	return strings.TrimSpace(tag), false
+}
+
+func processMember(field reflect.StructField, value reflect.Value,
+	collector map[string]map[string]interface{}) error {
+	var key string
+	var opt *fieldOptions
+	var err error
+	tag, ok := getTag(field)
+	if !ok {
+		tag = emptyTag
+		key = field.Name
+	} else {
+		key, opt, err = parseKeyAndOptions(tag, field)
+		if err != nil {
+			return err
+		}
+
+		if err = validate(field, value, opt); err != nil {
+			return err
+		}
+	}
+
+	val := value.Interface()
+	if opt != nil && opt.FromString {
+		val = fmt.Sprint(val)
+	}
+
+	m, ok := collector[tag]
+	if ok {
+		m[key] = val
+	} else {
+		m = map[string]interface{}{
+			key: val,
+		}
+	}
+	collector[tag] = m
+
+	return nil
+}
+
+func validate(field reflect.StructField, value reflect.Value, opt *fieldOptions) error {
+	if opt == nil || !opt.Optional {
+		if err := validateOptional(field, value); err != nil {
+			return err
+		}
+	}
+
+	if opt == nil {
+		return nil
+	}
+
+	if opt.Optional && value.IsZero() {
+		return nil
+	}
+
+	if len(opt.Options) > 0 {
+		if err := validateOptions(value, opt); err != nil {
+			return err
+		}
+	}
+
+	if opt.Range != nil {
+		if err := validateRange(value, opt); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func validateOptional(field reflect.StructField, value reflect.Value) error {
+	switch field.Type.Kind() {
+	case reflect.Ptr:
+		if value.IsNil() {
+			return fmt.Errorf("field %q is nil", field.Name)
+		}
+	case reflect.Array, reflect.Slice, reflect.Map:
+		if value.IsNil() || value.Len() == 0 {
+			return fmt.Errorf("field %q is empty", field.Name)
+		}
+	}
+
+	return nil
+}
+
+func validateOptions(value reflect.Value, opt *fieldOptions) error {
+	var found bool
+	val := fmt.Sprint(value.Interface())
+	for i := range opt.Options {
+		if opt.Options[i] == val {
+			found = true
+			break
+		}
+	}
+	if !found {
+		return fmt.Errorf("field %q not in options", val)
+	}
+
+	return nil
+}
+
+func validateRange(value reflect.Value, opt *fieldOptions) error {
+	var val float64
+	switch v := value.Interface().(type) {
+	case int:
+		val = float64(v)
+	case int8:
+		val = float64(v)
+	case int16:
+		val = float64(v)
+	case int32:
+		val = float64(v)
+	case int64:
+		val = float64(v)
+	case uint:
+		val = float64(v)
+	case uint8:
+		val = float64(v)
+	case uint16:
+		val = float64(v)
+	case uint32:
+		val = float64(v)
+	case uint64:
+		val = float64(v)
+	case float32:
+		val = float64(v)
+	case float64:
+		val = v
+	default:
+		return fmt.Errorf("unknown support type for range %q", value.Type().String())
+	}
+
+	// validates [left, right], [left, right), (left, right], (left, right)
+	if val < opt.Range.left ||
+		(!opt.Range.leftInclude && val == opt.Range.left) ||
+		val > opt.Range.right ||
+		(!opt.Range.rightInclude && val == opt.Range.right) {
+		return fmt.Errorf("%v out of range", value.Interface())
+	}
+
+	return nil
+}

core/mapping/marshaler_test.go

@@ -0,0 +1,274 @@
+package mapping
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestMarshal(t *testing.T) {
+	v := struct {
+		Name      string `path:"name"`
+		Address   string `json:"address,options=[beijing,shanghai]"`
+		Age       int    `json:"age"`
+		Anonymous bool
+	}{
+		Name:      "kevin",
+		Address:   "shanghai",
+		Age:       20,
+		Anonymous: true,
+	}
+
+	m, err := Marshal(v)
+	assert.Nil(t, err)
+	assert.Equal(t, "kevin", m["path"]["name"])
+	assert.Equal(t, "shanghai", m["json"]["address"])
+	assert.Equal(t, 20, m["json"]["age"].(int))
+	assert.True(t, m[emptyTag]["Anonymous"].(bool))
+}
+
+func TestMarshal_Ptr(t *testing.T) {
+	v := &struct {
+		Name      string `path:"name"`
+		Address   string `json:"address,options=[beijing,shanghai]"`
+		Age       int    `json:"age"`
+		Anonymous bool
+	}{
+		Name:      "kevin",
+		Address:   "shanghai",
+		Age:       20,
+		Anonymous: true,
+	}
+
+	m, err := Marshal(v)
+	assert.Nil(t, err)
+	assert.Equal(t, "kevin", m["path"]["name"])
+	assert.Equal(t, "shanghai", m["json"]["address"])
+	assert.Equal(t, 20, m["json"]["age"].(int))
+	assert.True(t, m[emptyTag]["Anonymous"].(bool))
+}
+
+func TestMarshal_OptionalPtr(t *testing.T) {
+	var val = 1
+	v := struct {
+		Age *int `json:"age"`
+	}{
+		Age: &val,
+	}
+
+	m, err := Marshal(v)
+	assert.Nil(t, err)
+	assert.Equal(t, 1, *m["json"]["age"].(*int))
+}
+
+func TestMarshal_OptionalPtrNil(t *testing.T) {
+	v := struct {
+		Age *int `json:"age"`
+	}{}
+
+	_, err := Marshal(v)
+	assert.NotNil(t, err)
+}
+
+func TestMarshal_BadOptions(t *testing.T) {
+	v := struct {
+		Name string `json:"name,options"`
+	}{
+		Name: "kevin",
+	}
+
+	_, err := Marshal(v)
+	assert.NotNil(t, err)
+}
+
+func TestMarshal_NotInOptions(t *testing.T) {
+	v := struct {
+		Name string `json:"name,options=[a,b]"`
+	}{
+		Name: "kevin",
+	}
+
+	_, err := Marshal(v)
+	assert.NotNil(t, err)
+}
+
+func TestMarshal_NotInOptionsOptional(t *testing.T) {
+	v := struct {
+		Name string `json:"name,options=[a,b],optional"`
+	}{}
+
+	_, err := Marshal(v)
+	assert.Nil(t, err)
+}
+
+func TestMarshal_NotInOptionsOptionalWrongValue(t *testing.T) {
+	v := struct {
+		Name string `json:"name,options=[a,b],optional"`
+	}{
+		Name: "kevin",
+	}
+
+	_, err := Marshal(v)
+	assert.NotNil(t, err)
+}
+
+func TestMarshal_Nested(t *testing.T) {
+	type address struct {
+		Country string `json:"country"`
+		City    string `json:"city"`
+	}
+	v := struct {
+		Name    string  `json:"name,options=[kevin,wan]"`
+		Address address `json:"address"`
+	}{
+		Name: "kevin",
+		Address: address{
+			Country: "China",
+			City:    "Shanghai",
+		},
+	}
+
+	m, err := Marshal(v)
+	assert.Nil(t, err)
+	assert.Equal(t, "kevin", m["json"]["name"])
+	assert.Equal(t, "China", m["json"]["address"].(address).Country)
+	assert.Equal(t, "Shanghai", m["json"]["address"].(address).City)
+}
+
+func TestMarshal_NestedPtr(t *testing.T) {
+	type address struct {
+		Country string `json:"country"`
+		City    string `json:"city"`
+	}
+	v := struct {
+		Name    string   `json:"name,options=[kevin,wan]"`
+		Address *address `json:"address"`
+	}{
+		Name: "kevin",
+		Address: &address{
+			Country: "China",
+			City:    "Shanghai",
+		},
+	}
+
+	m, err := Marshal(v)
+	assert.Nil(t, err)
+	assert.Equal(t, "kevin", m["json"]["name"])
+	assert.Equal(t, "China", m["json"]["address"].(*address).Country)
+	assert.Equal(t, "Shanghai", m["json"]["address"].(*address).City)
+}
+
+func TestMarshal_Slice(t *testing.T) {
+	v := struct {
+		Name []string `json:"name"`
+	}{
+		Name: []string{"kevin", "wan"},
+	}
+
+	m, err := Marshal(v)
+	assert.Nil(t, err)
+	assert.ElementsMatch(t, []string{"kevin", "wan"}, m["json"]["name"].([]string))
+}
+
+func TestMarshal_SliceNil(t *testing.T) {
+	v := struct {
+		Name []string `json:"name"`
+	}{
+		Name: nil,
+	}
+
+	_, err := Marshal(v)
+	assert.NotNil(t, err)
+}
+
+func TestMarshal_Range(t *testing.T) {
+	v := struct {
+		Int     int     `json:"int,range=[1:3]"`
+		Int8    int8    `json:"int8,range=[1:3)"`
+		Int16   int16   `json:"int16,range=(1:3]"`
+		Int32   int32   `json:"int32,range=(1:3)"`
+		Int64   int64   `json:"int64,range=(1:3)"`
+		Uint    uint    `json:"uint,range=[1:3]"`
+		Uint8   uint8   `json:"uint8,range=[1:3)"`
+		Uint16  uint16  `json:"uint16,range=(1:3]"`
+		Uint32  uint32  `json:"uint32,range=(1:3)"`
+		Uint64  uint64  `json:"uint64,range=(1:3)"`
+		Float32 float32 `json:"float32,range=(1:3)"`
+		Float64 float64 `json:"float64,range=(1:3)"`
+	}{
+		Int:     1,
+		Int8:    1,
+		Int16:   2,
+		Int32:   2,
+		Int64:   2,
+		Uint:    1,
+		Uint8:   1,
+		Uint16:  2,
+		Uint32:  2,
+		Uint64:  2,
+		Float32: 2,
+		Float64: 2,
+	}
+
+	m, err := Marshal(v)
+	assert.Nil(t, err)
+	assert.Equal(t, 1, m["json"]["int"].(int))
+	assert.Equal(t, int8(1), m["json"]["int8"].(int8))
+	assert.Equal(t, int16(2), m["json"]["int16"].(int16))
+	assert.Equal(t, int32(2), m["json"]["int32"].(int32))
+	assert.Equal(t, int64(2), m["json"]["int64"].(int64))
+	assert.Equal(t, uint(1), m["json"]["uint"].(uint))
+	assert.Equal(t, uint8(1), m["json"]["uint8"].(uint8))
+	assert.Equal(t, uint16(2), m["json"]["uint16"].(uint16))
+	assert.Equal(t, uint32(2), m["json"]["uint32"].(uint32))
+	assert.Equal(t, uint64(2), m["json"]["uint64"].(uint64))
+	assert.Equal(t, float32(2), m["json"]["float32"].(float32))
+	assert.Equal(t, float64(2), m["json"]["float64"].(float64))
+}
+
+func TestMarshal_RangeOut(t *testing.T) {
+	tests := []interface{}{
+		struct {
+			Int int `json:"int,range=[1:3]"`
+		}{
+			Int: 4,
+		},
+		struct {
+			Int int `json:"int,range=(1:3]"`
+		}{
+			Int: 1,
+		},
+		struct {
+			Int int `json:"int,range=[1:3)"`
+		}{
+			Int: 3,
+		},
+		struct {
+			Int int `json:"int,range=(1:3)"`
+		}{
+			Int: 3,
+		},
+		struct {
+			Bool bool `json:"bool,range=(1:3)"`
+		}{
+			Bool: true,
+		},
+	}
+
+	for _, test := range tests {
+		_, err := Marshal(test)
+		assert.NotNil(t, err)
+	}
+}
+
+func TestMarshal_FromString(t *testing.T) {
+	v := struct {
+		Age int `json:"age,string"`
+	}{
+		Age: 10,
+	}
+
+	m, err := Marshal(v)
+	assert.Nil(t, err)
+	assert.Equal(t, "10", m["json"]["age"].(string))
+}

core/mapping/unmarshaler.go

@@ -97,10 +97,6 @@ func (u *Unmarshaler) unmarshalWithFullName(m Valuer, v interface{}, fullName st
 	numFields := rte.NumField()
 	for i := 0; i < numFields; i++ {
 		field := rte.Field(i)
-		if usingDifferentKeys(u.key, field) {
-			continue
-		}
-
 		if err := u.processField(field, rve.Field(i), m, fullName); err != nil {
 			return err
 		}
@@ -275,6 +271,10 @@ func (u *Unmarshaler) processFieldPrimitiveWithJSONNumber(field reflect.StructFi
 			return err
 		}
 
+		if iValue < 0 {
+			return fmt.Errorf("unmarshal %q with bad value %q", fullName, v.String())
+		}
+
 		value.SetUint(uint64(iValue))
 	case reflect.Float32, reflect.Float64:
 		fValue, err := v.Float64()
@@ -448,7 +448,15 @@ func (u *Unmarshaler) fillSlice(fieldType reflect.Type, value reflect.Value, map
 	dereffedBaseType := Deref(baseType)
 	dereffedBaseKind := dereffedBaseType.Kind()
 	refValue := reflect.ValueOf(mapValue)
+	if refValue.IsNil() {
+		return nil
+	}
+
 	conv := reflect.MakeSlice(reflect.SliceOf(baseType), refValue.Len(), refValue.Cap())
+	if refValue.Len() == 0 {
+		value.Set(conv)
+		return nil
+	}
 
 	var valid bool
 	for i := 0; i < refValue.Len(); i++ {
@@ -723,10 +731,10 @@ func fillWithSameType(field reflect.StructField, value reflect.Value, mapValue i
 	if field.Type.Kind() == reflect.Ptr {
 		baseType := Deref(field.Type)
 		target := reflect.New(baseType).Elem()
-		target.Set(reflect.ValueOf(mapValue))
+		setSameKindValue(baseType, target, mapValue)
 		value.Set(target.Addr())
 	} else {
-		value.Set(reflect.ValueOf(mapValue))
+		setSameKindValue(field.Type, value, mapValue)
 	}
 
 	return nil
@@ -801,3 +809,11 @@ func readKeys(key string) []string {
 
 	return keys
 }
+
+func setSameKindValue(targetType reflect.Type, target reflect.Value, value interface{}) {
+	if reflect.ValueOf(value).Type().AssignableTo(targetType) {
+		target.Set(reflect.ValueOf(value))
+	} else {
+		target.Set(reflect.ValueOf(value).Convert(targetType))
+	}
+}

core/mapping/unmarshaler_test.go

@@ -198,6 +198,49 @@ func TestUnmarshalIntWithDefault(t *testing.T) {
 	assert.Equal(t, 1, in.Int)
 }
 
+func TestUnmarshalBoolSliceRequired(t *testing.T) {
+	type inner struct {
+		Bools []bool `key:"bools"`
+	}
+
+	var in inner
+	assert.NotNil(t, UnmarshalKey(map[string]interface{}{}, &in))
+}
+
+func TestUnmarshalBoolSliceNil(t *testing.T) {
+	type inner struct {
+		Bools []bool `key:"bools,optional"`
+	}
+
+	var in inner
+	assert.Nil(t, UnmarshalKey(map[string]interface{}{}, &in))
+	assert.Nil(t, in.Bools)
+}
+
+func TestUnmarshalBoolSliceNilExplicit(t *testing.T) {
+	type inner struct {
+		Bools []bool `key:"bools,optional"`
+	}
+
+	var in inner
+	assert.Nil(t, UnmarshalKey(map[string]interface{}{
+		"bools": nil,
+	}, &in))
+	assert.Nil(t, in.Bools)
+}
+
+func TestUnmarshalBoolSliceEmpty(t *testing.T) {
+	type inner struct {
+		Bools []bool `key:"bools,optional"`
+	}
+
+	var in inner
+	assert.Nil(t, UnmarshalKey(map[string]interface{}{
+		"bools": []bool{},
+	}, &in))
+	assert.Empty(t, in.Bools)
+}
+
 func TestUnmarshalBoolSliceWithDefault(t *testing.T) {
 	type inner struct {
 		Bools []bool `key:"bools,default=[true,false]"`
@@ -330,28 +373,34 @@ func TestUnmarshalFloat(t *testing.T) {
 
 func TestUnmarshalInt64Slice(t *testing.T) {
 	var v struct {
-		Ages []int64 `key:"ages"`
+		Ages  []int64 `key:"ages"`
+		Slice []int64 `key:"slice"`
 	}
 	m := map[string]interface{}{
-		"ages": []int64{1, 2},
+		"ages":  []int64{1, 2},
+		"slice": []interface{}{},
 	}
 
 	ast := assert.New(t)
 	ast.Nil(UnmarshalKey(m, &v))
 	ast.ElementsMatch([]int64{1, 2}, v.Ages)
+	ast.Equal([]int64{}, v.Slice)
 }
 
 func TestUnmarshalIntSlice(t *testing.T) {
 	var v struct {
-		Ages []int `key:"ages"`
+		Ages  []int `key:"ages"`
+		Slice []int `key:"slice"`
 	}
 	m := map[string]interface{}{
-		"ages": []int{1, 2},
+		"ages":  []int{1, 2},
+		"slice": []interface{}{},
 	}
 
 	ast := assert.New(t)
 	ast.Nil(UnmarshalKey(m, &v))
 	ast.ElementsMatch([]int{1, 2}, v.Ages)
+	ast.Equal([]int{}, v.Slice)
 }
 
 func TestUnmarshalString(t *testing.T) {
@@ -938,6 +987,43 @@ func TestUnmarshalWithStringOptionsCorrect(t *testing.T) {
 	ast.Equal("2", in.Correct)
 }
 
+func TestUnmarshalOptionsOptional(t *testing.T) {
+	type inner struct {
+		Value         string `key:"value,options=first|second,optional"`
+		OptionalValue string `key:"optional_value,options=first|second,optional"`
+		Foo           string `key:"foo,options=[bar,baz]"`
+		Correct       string `key:"correct,options=1|2"`
+	}
+	m := map[string]interface{}{
+		"value":   "first",
+		"foo":     "bar",
+		"correct": "2",
+	}
+
+	var in inner
+	ast := assert.New(t)
+	ast.Nil(UnmarshalKey(m, &in))
+	ast.Equal("first", in.Value)
+	ast.Equal("", in.OptionalValue)
+	ast.Equal("bar", in.Foo)
+	ast.Equal("2", in.Correct)
+}
+
+func TestUnmarshalOptionsOptionalWrongValue(t *testing.T) {
+	type inner struct {
+		Value         string `key:"value,options=first|second,optional"`
+		OptionalValue string `key:"optional_value,options=first|second,optional"`
+		WrongValue    string `key:"wrong_value,options=first|second,optional"`
+	}
+	m := map[string]interface{}{
+		"value":       "first",
+		"wrong_value": "third",
+	}
+
+	var in inner
+	assert.NotNil(t, UnmarshalKey(m, &in))
+}
+
 func TestUnmarshalStringOptionsWithStringOptionsNotString(t *testing.T) {
 	type inner struct {
 		Value   string `key:"value,options=first|second"`
@@ -2611,6 +2697,86 @@ func TestUnmarshalJsonWithoutKey(t *testing.T) {
 	assert.Equal(t, "2", res.B)
 }
 
+func TestUnmarshalJsonUintNegative(t *testing.T) {
+	payload := `{"a": -1}`
+	var res struct {
+		A uint `json:"a"`
+	}
+	reader := strings.NewReader(payload)
+	err := UnmarshalJsonReader(reader, &res)
+	assert.NotNil(t, err)
+}
+
+func TestUnmarshalJsonDefinedInt(t *testing.T) {
+	type Int int
+	var res struct {
+		A Int `json:"a"`
+	}
+	payload := `{"a": -1}`
+	reader := strings.NewReader(payload)
+	err := UnmarshalJsonReader(reader, &res)
+	assert.Nil(t, err)
+	assert.Equal(t, Int(-1), res.A)
+}
+
+func TestUnmarshalJsonDefinedString(t *testing.T) {
+	type String string
+	var res struct {
+		A String `json:"a"`
+	}
+	payload := `{"a": "foo"}`
+	reader := strings.NewReader(payload)
+	err := UnmarshalJsonReader(reader, &res)
+	assert.Nil(t, err)
+	assert.Equal(t, String("foo"), res.A)
+}
+
+func TestUnmarshalJsonDefinedStringPtr(t *testing.T) {
+	type String string
+	var res struct {
+		A *String `json:"a"`
+	}
+	payload := `{"a": "foo"}`
+	reader := strings.NewReader(payload)
+	err := UnmarshalJsonReader(reader, &res)
+	assert.Nil(t, err)
+	assert.Equal(t, String("foo"), *res.A)
+}
+
+func TestUnmarshalJsonReaderComplex(t *testing.T) {
+	type (
+		MyInt      int
+		MyTxt      string
+		MyTxtArray []string
+
+		Req struct {
+			MyInt      MyInt      `json:"my_int"` // int.. ok
+			MyTxtArray MyTxtArray `json:"my_txt_array"`
+			MyTxt      MyTxt      `json:"my_txt"` // but string is not assignable
+			Int        int        `json:"int"`
+			Txt        string     `json:"txt"`
+		}
+	)
+	body := `{
+  "my_int": 100,
+  "my_txt_array": [
+    "a",
+    "b"
+  ],
+  "my_txt": "my_txt",
+  "int": 200,
+  "txt": "txt"
+}`
+	var req Req
+	err := UnmarshalJsonReader(strings.NewReader(body), &req)
+	assert.Nil(t, err)
+	assert.Equal(t, MyInt(100), req.MyInt)
+	assert.Equal(t, MyTxt("my_txt"), req.MyTxt)
+	assert.EqualValues(t, MyTxtArray([]string{"a", "b"}), req.MyTxtArray)
+	assert.Equal(t, 200, req.Int)
+	assert.Equal(t, "txt", req.Txt)
+}
+
 func BenchmarkDefaultValue(b *testing.B) {
 	for i := 0; i < b.N; i++ {
 		var a struct {

core/prof/runtime.go

@@ -11,6 +11,7 @@ const (
 	mega            = 1024 * 1024
 )
 
+// DisplayStats prints the goroutine, memory, GC stats with given interval, default to 5 seconds.
 func DisplayStats(interval ...time.Duration) {
 	duration := defaultInterval
 	for _, val := range interval {

core/stat/remotewriter.go

@@ -10,7 +10,10 @@ import (
 	"github.com/zeromicro/go-zero/core/logx"
 )
 
-const httpTimeout = time.Second * 5
+const (
+	httpTimeout     = time.Second * 5
+	jsonContentType = "application/json; charset=utf-8"
+)
 
 // ErrWriteFailed is an error that indicates failed to submit a StatReport.
 var ErrWriteFailed = errors.New("submit failed")
@@ -36,7 +39,7 @@ func (rw *RemoteWriter) Write(report *StatReport) error {
 	client := &http.Client{
 		Timeout: httpTimeout,
 	}
-	resp, err := client.Post(rw.endpoint, "application/json", bytes.NewBuffer(bs))
+	resp, err := client.Post(rw.endpoint, jsonContentType, bytes.NewReader(bs))
 	if err != nil {
 		return err
 	}

core/stores/builder/builder.go

@@ -41,6 +41,16 @@ func RawFieldNames(in interface{}, postgresSql ...bool) []string {
 				out = append(out, fmt.Sprintf("`%s`", fi.Name))
 			}
 		default:
+			// get tag name with the tag opton, e.g.:
+			// `db:"id"`
+			// `db:"id,type=char,length=16"`
+			// `db:",type=char,length=16"`
+			if strings.Contains(tagv, ",") {
+				tagv = strings.TrimSpace(strings.Split(tagv, ",")[0])
+			}
+			if len(tagv) == 0 {
+				tagv = fi.Name
+			}
 			if pg {
 				out = append(out, tagv)
 			} else {

core/stores/builder/builder_test.go

@@ -22,3 +22,20 @@ func TestFieldNames(t *testing.T) {
 		assert.Equal(t, expected, out)
 	})
 }
+
+type mockedUserWithOptions struct {
+	ID       string `db:"id" json:"id,omitempty"`
+	UserName string `db:"user_name,type=varchar,length=255" json:"userName,omitempty"`
+	Sex      int    `db:"sex" json:"sex,omitempty"`
+	UUID     string `db:",type=varchar,length=16" uuid:"uuid,omitempty"`
+	Age      int    `db:"age" json:"age"`
+}
+
+func TestFieldNamesWithTagOptions(t *testing.T) {
+	t.Run("new", func(t *testing.T) {
+		var u mockedUserWithOptions
+		out := RawFieldNames(&u)
+		expected := []string{"`id`", "`user_name`", "`sex`", "`UUID`", "`age`"}
+		assert.Equal(t, expected, out)
+	})
+}

core/stores/kv/store_test.go

@@ -490,6 +490,29 @@ func TestRedis_SetExNx(t *testing.T) {
 	})
 }
 
+func TestRedis_Getset(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.GetSet("hello", "world")
+	assert.NotNil(t, err)
+
+	runOnCluster(t, func(client Store) {
+		val, err := client.GetSet("hello", "world")
+		assert.Nil(t, err)
+		assert.Equal(t, "", val)
+		val, err = client.Get("hello")
+		assert.Nil(t, err)
+		assert.Equal(t, "world", val)
+		val, err = client.GetSet("hello", "newworld")
+		assert.Nil(t, err)
+		assert.Equal(t, "world", val)
+		val, err = client.Get("hello")
+		assert.Nil(t, err)
+		assert.Equal(t, "newworld", val)
+		_, err = client.Del("hello")
+		assert.Nil(t, err)
+	})
+}
+
 func TestRedis_SetGetDelHashField(t *testing.T) {
 	store := clusterStore{dispatcher: hash.NewConsistentHash()}
 	err := store.Hset("key", "field", "value")

core/stores/mon/bulkinserter.go

@@ -0,0 +1,91 @@
+package mon
+
+import (
+	"context"
+	"time"
+
+	"github.com/zeromicro/go-zero/core/executors"
+	"github.com/zeromicro/go-zero/core/logx"
+	"go.mongodb.org/mongo-driver/mongo"
+)
+
+const (
+	flushInterval = time.Second
+	maxBulkRows   = 1000
+)
+
+type (
+	// ResultHandler is a handler that used to handle results.
+	ResultHandler func(*mongo.InsertManyResult, error)
+
+	// A BulkInserter is used to insert bulk of mongo records.
+	BulkInserter struct {
+		executor *executors.PeriodicalExecutor
+		inserter *dbInserter
+	}
+)
+
+// NewBulkInserter returns a BulkInserter.
+func NewBulkInserter(coll *mongo.Collection, interval ...time.Duration) *BulkInserter {
+	inserter := &dbInserter{
+		collection: coll,
+	}
+
+	duration := flushInterval
+	if len(interval) > 0 {
+		duration = interval[0]
+	}
+
+	return &BulkInserter{
+		executor: executors.NewPeriodicalExecutor(duration, inserter),
+		inserter: inserter,
+	}
+}
+
+// Flush flushes the inserter, writes all pending records.
+func (bi *BulkInserter) Flush() {
+	bi.executor.Flush()
+}
+
+// Insert inserts doc.
+func (bi *BulkInserter) Insert(doc interface{}) {
+	bi.executor.Add(doc)
+}
+
+// SetResultHandler sets the result handler.
+func (bi *BulkInserter) SetResultHandler(handler ResultHandler) {
+	bi.executor.Sync(func() {
+		bi.inserter.resultHandler = handler
+	})
+}
+
+type dbInserter struct {
+	collection    *mongo.Collection
+	documents     []interface{}
+	resultHandler ResultHandler
+}
+
+func (in *dbInserter) AddTask(doc interface{}) bool {
+	in.documents = append(in.documents, doc)
+	return len(in.documents) >= maxBulkRows
+}
+
+func (in *dbInserter) Execute(objs interface{}) {
+	docs := objs.([]interface{})
+	if len(docs) == 0 {
+		return
+	}
+
+	result, err := in.collection.InsertMany(context.Background(), docs)
+	if in.resultHandler != nil {
+		in.resultHandler(result, err)
+	} else if err != nil {
+		logx.Error(err)
+	}
+}
+
+func (in *dbInserter) RemoveAll() interface{} {
+	documents := in.documents
+	in.documents = nil
+	return documents
+}

core/stores/mon/bulkinserter_test.go

@@ -0,0 +1,27 @@
+package mon
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/mongo"
+	"go.mongodb.org/mongo-driver/mongo/integration/mtest"
+)
+
+func TestBulkInserter(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "ok", Value: 1}}...))
+		bulk := NewBulkInserter(mt.Coll)
+		bulk.SetResultHandler(func(result *mongo.InsertManyResult, err error) {
+			assert.Nil(t, err)
+			assert.Equal(t, 2, len(result.InsertedIDs))
+		})
+		bulk.Insert(bson.D{{Key: "foo", Value: "bar"}})
+		bulk.Insert(bson.D{{Key: "foo", Value: "baz"}})
+		bulk.Flush()
+	})
+}

core/stores/mon/clientmanager.go

@@ -0,0 +1,51 @@
+package mon
+
+import (
+	"context"
+	"io"
+	"time"
+
+	"github.com/zeromicro/go-zero/core/syncx"
+	"go.mongodb.org/mongo-driver/mongo"
+	mopt "go.mongodb.org/mongo-driver/mongo/options"
+)
+
+const defaultTimeout = time.Second
+
+var clientManager = syncx.NewResourceManager()
+
+// ClosableClient wraps *mongo.Client and provides a Close method.
+type ClosableClient struct {
+	*mongo.Client
+}
+
+// Close disconnects the underlying *mongo.Client.
+func (cs *ClosableClient) Close() error {
+	return cs.Client.Disconnect(context.Background())
+}
+
+// Inject injects a *mongo.Client into the client manager.
+// Typically, this is used to inject a *mongo.Client for test purpose.
+func Inject(key string, client *mongo.Client) {
+	clientManager.Inject(key, &ClosableClient{client})
+}
+
+func getClient(url string) (*mongo.Client, error) {
+	val, err := clientManager.GetResource(url, func() (io.Closer, error) {
+		cli, err := mongo.Connect(context.Background(), mopt.Client().ApplyURI(url))
+		if err != nil {
+			return nil, err
+		}
+
+		concurrentSess := &ClosableClient{
+			Client: cli,
+		}
+
+		return concurrentSess, nil
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	return val.(*ClosableClient).Client, nil
+}

core/stores/mon/clientmanager_test.go

@@ -0,0 +1,20 @@
+package mon
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"go.mongodb.org/mongo-driver/mongo/integration/mtest"
+)
+
+func TestClientManger_getClient(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		Inject(mtest.ClusterURI(), mt.Client)
+		cli, err := getClient(mtest.ClusterURI())
+		assert.Nil(t, err)
+		assert.Equal(t, mt.Client, cli)
+	})
+}

core/stores/mon/collection.go

@@ -0,0 +1,490 @@
+package mon
+
+import (
+	"context"
+	"encoding/json"
+	"time"
+
+	"github.com/zeromicro/go-zero/core/breaker"
+	"github.com/zeromicro/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/core/timex"
+	"github.com/zeromicro/go-zero/core/trace"
+	"go.mongodb.org/mongo-driver/mongo"
+	mopt "go.mongodb.org/mongo-driver/mongo/options"
+	"go.mongodb.org/mongo-driver/x/mongo/driver/session"
+	"go.opentelemetry.io/otel"
+	tracesdk "go.opentelemetry.io/otel/trace"
+)
+
+const (
+	defaultSlowThreshold = time.Millisecond * 500
+	// spanName is the span name of the mongo calls.
+	spanName = "mongo"
+)
+
+// ErrNotFound is an alias of mongo.ErrNoDocuments
+var ErrNotFound = mongo.ErrNoDocuments
+
+type (
+	// Collection defines a MongoDB collection.
+	Collection interface {
+		// Aggregate executes an aggregation pipeline.
+		Aggregate(ctx context.Context, pipeline interface{}, opts ...*mopt.AggregateOptions) (
+			*mongo.Cursor, error)
+		// BulkWrite performs a bulk write operation.
+		BulkWrite(ctx context.Context, models []mongo.WriteModel, opts ...*mopt.BulkWriteOptions) (
+			*mongo.BulkWriteResult, error)
+		// Clone creates a copy of this collection with the same settings.
+		Clone(opts ...*mopt.CollectionOptions) (*mongo.Collection, error)
+		// CountDocuments returns the number of documents in the collection that match the filter.
+		CountDocuments(ctx context.Context, filter interface{}, opts ...*mopt.CountOptions) (int64, error)
+		// Database returns the database that this collection is a part of.
+		Database() *mongo.Database
+		// DeleteMany deletes documents from the collection that match the filter.
+		DeleteMany(ctx context.Context, filter interface{}, opts ...*mopt.DeleteOptions) (
+			*mongo.DeleteResult, error)
+		// DeleteOne deletes at most one document from the collection that matches the filter.
+		DeleteOne(ctx context.Context, filter interface{}, opts ...*mopt.DeleteOptions) (
+			*mongo.DeleteResult, error)
+		// Distinct returns a list of distinct values for the given key across the collection.
+		Distinct(ctx context.Context, fieldName string, filter interface{},
+			opts ...*mopt.DistinctOptions) ([]interface{}, error)
+		// Drop drops this collection from database.
+		Drop(ctx context.Context) error
+		// EstimatedDocumentCount returns an estimate of the count of documents in a collection
+		// using collection metadata.
+		EstimatedDocumentCount(ctx context.Context, opts ...*mopt.EstimatedDocumentCountOptions) (int64, error)
+		// Find finds the documents matching the provided filter.
+		Find(ctx context.Context, filter interface{}, opts ...*mopt.FindOptions) (*mongo.Cursor, error)
+		// FindOne returns up to one document that matches the provided filter.
+		FindOne(ctx context.Context, filter interface{}, opts ...*mopt.FindOneOptions) (
+			*mongo.SingleResult, error)
+		// FindOneAndDelete returns at most one document that matches the filter. If the filter
+		// matches multiple documents, only the first document is deleted.
+		FindOneAndDelete(ctx context.Context, filter interface{}, opts ...*mopt.FindOneAndDeleteOptions) (
+			*mongo.SingleResult, error)
+		// FindOneAndReplace returns at most one document that matches the filter. If the filter
+		// matches multiple documents, FindOneAndReplace returns the first document in the
+		// collection that matches the filter.
+		FindOneAndReplace(ctx context.Context, filter interface{}, replacement interface{},
+			opts ...*mopt.FindOneAndReplaceOptions) (*mongo.SingleResult, error)
+		// FindOneAndUpdate returns at most one document that matches the filter. If the filter
+		// matches multiple documents, FindOneAndUpdate returns the first document in the
+		// collection that matches the filter.
+		FindOneAndUpdate(ctx context.Context, filter interface{}, update interface{},
+			opts ...*mopt.FindOneAndUpdateOptions) (*mongo.SingleResult, error)
+		// Indexes returns the index view for this collection.
+		Indexes() mongo.IndexView
+		// InsertMany inserts the provided documents.
+		InsertMany(ctx context.Context, documents []interface{}, opts ...*mopt.InsertManyOptions) (
+			*mongo.InsertManyResult, error)
+		// InsertOne inserts the provided document.
+		InsertOne(ctx context.Context, document interface{}, opts ...*mopt.InsertOneOptions) (
+			*mongo.InsertOneResult, error)
+		// ReplaceOne replaces at most one document that matches the filter.
+		ReplaceOne(ctx context.Context, filter interface{}, replacement interface{},
+			opts ...*mopt.ReplaceOptions) (*mongo.UpdateResult, error)
+		// UpdateByID updates a single document matching the provided filter.
+		UpdateByID(ctx context.Context, id interface{}, update interface{},
+			opts ...*mopt.UpdateOptions) (*mongo.UpdateResult, error)
+		// UpdateMany updates the provided documents.
+		UpdateMany(ctx context.Context, filter interface{}, update interface{},
+			opts ...*mopt.UpdateOptions) (*mongo.UpdateResult, error)
+		// UpdateOne updates a single document matching the provided filter.
+		UpdateOne(ctx context.Context, filter interface{}, update interface{},
+			opts ...*mopt.UpdateOptions) (*mongo.UpdateResult, error)
+		// Watch returns a change stream cursor used to receive notifications of changes to the collection.
+		Watch(ctx context.Context, pipeline interface{}, opts ...*mopt.ChangeStreamOptions) (
+			*mongo.ChangeStream, error)
+	}
+
+	decoratedCollection struct {
+		*mongo.Collection
+		name string
+		brk  breaker.Breaker
+	}
+
+	keepablePromise struct {
+		promise breaker.Promise
+		log     func(error)
+	}
+)
+
+func newCollection(collection *mongo.Collection, brk breaker.Breaker) Collection {
+	return &decoratedCollection{
+		Collection: collection,
+		name:       collection.Name(),
+		brk:        brk,
+	}
+}
+
+func (c *decoratedCollection) Aggregate(ctx context.Context, pipeline interface{},
+	opts ...*mopt.AggregateOptions) (cur *mongo.Cursor, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
+	err = c.brk.DoWithAcceptable(func() error {
+		starTime := timex.Now()
+		defer func() {
+			c.logDurationSimple("Aggregate", starTime, err)
+		}()
+
+		cur, err = c.Collection.Aggregate(ctx, pipeline, opts...)
+		return err
+	}, acceptable)
+	return
+}
+
+func (c *decoratedCollection) BulkWrite(ctx context.Context, models []mongo.WriteModel,
+	opts ...*mopt.BulkWriteOptions) (res *mongo.BulkWriteResult, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
+	err = c.brk.DoWithAcceptable(func() error {
+		startTime := timex.Now()
+		defer func() {
+			c.logDurationSimple("BulkWrite", startTime, err)
+		}()
+
+		res, err = c.Collection.BulkWrite(ctx, models, opts...)
+		return err
+	}, acceptable)
+	return
+}
+
+func (c *decoratedCollection) CountDocuments(ctx context.Context, filter interface{},
+	opts ...*mopt.CountOptions) (count int64, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
+	err = c.brk.DoWithAcceptable(func() error {
+		startTime := timex.Now()
+		defer func() {
+			c.logDurationSimple("CountDocuments", startTime, err)
+		}()
+
+		count, err = c.Collection.CountDocuments(ctx, filter, opts...)
+		return err
+	}, acceptable)
+	return
+}
+
+func (c *decoratedCollection) DeleteMany(ctx context.Context, filter interface{},
+	opts ...*mopt.DeleteOptions) (res *mongo.DeleteResult, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
+	err = c.brk.DoWithAcceptable(func() error {
+		startTime := timex.Now()
+		defer func() {
+			c.logDurationSimple("DeleteMany", startTime, err)
+		}()
+
+		res, err = c.Collection.DeleteMany(ctx, filter, opts...)
+		return err
+	}, acceptable)
+	return
+}
+
+func (c *decoratedCollection) DeleteOne(ctx context.Context, filter interface{},
+	opts ...*mopt.DeleteOptions) (res *mongo.DeleteResult, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
+	err = c.brk.DoWithAcceptable(func() error {
+		startTime := timex.Now()
+		defer func() {
+			c.logDuration("DeleteOne", startTime, err, filter)
+		}()
+
+		res, err = c.Collection.DeleteOne(ctx, filter, opts...)
+		return err
+	}, acceptable)
+	return
+}
+
+func (c *decoratedCollection) Distinct(ctx context.Context, fieldName string, filter interface{},
+	opts ...*mopt.DistinctOptions) (val []interface{}, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
+	err = c.brk.DoWithAcceptable(func() error {
+		startTime := timex.Now()
+		defer func() {
+			c.logDurationSimple("Distinct", startTime, err)
+		}()
+
+		val, err = c.Collection.Distinct(ctx, fieldName, filter, opts...)
+		return err
+	}, acceptable)
+	return
+}
+
+func (c *decoratedCollection) EstimatedDocumentCount(ctx context.Context,
+	opts ...*mopt.EstimatedDocumentCountOptions) (val int64, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
+	err = c.brk.DoWithAcceptable(func() error {
+		startTime := timex.Now()
+		defer func() {
+			c.logDurationSimple("EstimatedDocumentCount", startTime, err)
+		}()
+
+		val, err = c.Collection.EstimatedDocumentCount(ctx, opts...)
+		return err
+	}, acceptable)
+	return
+}
+
+func (c *decoratedCollection) Find(ctx context.Context, filter interface{},
+	opts ...*mopt.FindOptions) (cur *mongo.Cursor, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
+	err = c.brk.DoWithAcceptable(func() error {
+		startTime := timex.Now()
+		defer func() {
+			c.logDuration("Find", startTime, err, filter)
+		}()
+
+		cur, err = c.Collection.Find(ctx, filter, opts...)
+		return err
+	}, acceptable)
+	return
+}
+
+func (c *decoratedCollection) FindOne(ctx context.Context, filter interface{},
+	opts ...*mopt.FindOneOptions) (res *mongo.SingleResult, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
+	err = c.brk.DoWithAcceptable(func() error {
+		startTime := timex.Now()
+		defer func() {
+			c.logDuration("FindOne", startTime, err, filter)
+		}()
+
+		res = c.Collection.FindOne(ctx, filter, opts...)
+		err = res.Err()
+		return err
+	}, acceptable)
+	return
+}
+
+func (c *decoratedCollection) FindOneAndDelete(ctx context.Context, filter interface{},
+	opts ...*mopt.FindOneAndDeleteOptions) (res *mongo.SingleResult, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
+	err = c.brk.DoWithAcceptable(func() error {
+		startTime := timex.Now()
+		defer func() {
+			c.logDuration("FindOneAndDelete", startTime, err, filter)
+		}()
+
+		res = c.Collection.FindOneAndDelete(ctx, filter, opts...)
+		err = res.Err()
+		return err
+	}, acceptable)
+	return
+}
+
+func (c *decoratedCollection) FindOneAndReplace(ctx context.Context, filter interface{},
+	replacement interface{}, opts ...*mopt.FindOneAndReplaceOptions) (
+	res *mongo.SingleResult, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
+	err = c.brk.DoWithAcceptable(func() error {
+		startTime := timex.Now()
+		defer func() {
+			c.logDuration("FindOneAndReplace", startTime, err, filter, replacement)
+		}()
+
+		res = c.Collection.FindOneAndReplace(ctx, filter, replacement, opts...)
+		err = res.Err()
+		return err
+	}, acceptable)
+	return
+}
+
+func (c *decoratedCollection) FindOneAndUpdate(ctx context.Context, filter interface{}, update interface{},
+	opts ...*mopt.FindOneAndUpdateOptions) (res *mongo.SingleResult, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
+	err = c.brk.DoWithAcceptable(func() error {
+		startTime := timex.Now()
+		defer func() {
+			c.logDuration("FindOneAndUpdate", startTime, err, filter, update)
+		}()
+
+		res = c.Collection.FindOneAndUpdate(ctx, filter, update, opts...)
+		err = res.Err()
+		return err
+	}, acceptable)
+	return
+}
+
+func (c *decoratedCollection) InsertMany(ctx context.Context, documents []interface{},
+	opts ...*mopt.InsertManyOptions) (res *mongo.InsertManyResult, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
+	err = c.brk.DoWithAcceptable(func() error {
+		startTime := timex.Now()
+		defer func() {
+			c.logDurationSimple("InsertMany", startTime, err)
+		}()
+
+		res, err = c.Collection.InsertMany(ctx, documents, opts...)
+		return err
+	}, acceptable)
+	return
+}
+
+func (c *decoratedCollection) InsertOne(ctx context.Context, document interface{},
+	opts ...*mopt.InsertOneOptions) (res *mongo.InsertOneResult, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
+	err = c.brk.DoWithAcceptable(func() error {
+		startTime := timex.Now()
+		defer func() {
+			c.logDuration("InsertOne", startTime, err, document)
+		}()
+
+		res, err = c.Collection.InsertOne(ctx, document, opts...)
+		return err
+	}, acceptable)
+	return
+}
+
+func (c *decoratedCollection) ReplaceOne(ctx context.Context, filter interface{}, replacement interface{},
+	opts ...*mopt.ReplaceOptions) (res *mongo.UpdateResult, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
+	err = c.brk.DoWithAcceptable(func() error {
+		startTime := timex.Now()
+		defer func() {
+			c.logDuration("ReplaceOne", startTime, err, filter, replacement)
+		}()
+
+		res, err = c.Collection.ReplaceOne(ctx, filter, replacement, opts...)
+		return err
+	}, acceptable)
+	return
+}
+
+func (c *decoratedCollection) UpdateByID(ctx context.Context, id interface{}, update interface{},
+	opts ...*mopt.UpdateOptions) (res *mongo.UpdateResult, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
+	err = c.brk.DoWithAcceptable(func() error {
+		startTime := timex.Now()
+		defer func() {
+			c.logDuration("UpdateByID", startTime, err, id, update)
+		}()
+
+		res, err = c.Collection.UpdateByID(ctx, id, update, opts...)
+		return err
+	}, acceptable)
+	return
+}
+
+func (c *decoratedCollection) UpdateMany(ctx context.Context, filter interface{}, update interface{},
+	opts ...*mopt.UpdateOptions) (res *mongo.UpdateResult, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
+	err = c.brk.DoWithAcceptable(func() error {
+		startTime := timex.Now()
+		defer func() {
+			c.logDurationSimple("UpdateMany", startTime, err)
+		}()
+
+		res, err = c.Collection.UpdateMany(ctx, filter, update, opts...)
+		return err
+	}, acceptable)
+	return
+}
+
+func (c *decoratedCollection) UpdateOne(ctx context.Context, filter interface{}, update interface{},
+	opts ...*mopt.UpdateOptions) (res *mongo.UpdateResult, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
+	err = c.brk.DoWithAcceptable(func() error {
+		startTime := timex.Now()
+		defer func() {
+			c.logDuration("UpdateOne", startTime, err, filter, update)
+		}()
+
+		res, err = c.Collection.UpdateOne(ctx, filter, update, opts...)
+		return err
+	}, acceptable)
+	return
+}
+
+func (c *decoratedCollection) logDuration(method string, startTime time.Duration, err error,
+	docs ...interface{}) {
+	duration := timex.Since(startTime)
+	content, e := json.Marshal(docs)
+	if e != nil {
+		logx.Error(err)
+	} else if err != nil {
+		if duration > slowThreshold.Load() {
+			logx.WithDuration(duration).Slowf("[MONGO] mongo(%s) - slowcall - %s - fail(%s) - %s",
+				c.name, method, err.Error(), string(content))
+		} else {
+			logx.WithDuration(duration).Infof("mongo(%s) - %s - fail(%s) - %s",
+				c.name, method, err.Error(), string(content))
+		}
+	} else {
+		if duration > slowThreshold.Load() {
+			logx.WithDuration(duration).Slowf("[MONGO] mongo(%s) - slowcall - %s - ok - %s",
+				c.name, method, string(content))
+		} else {
+			logx.WithDuration(duration).Infof("mongo(%s) - %s - ok - %s", c.name, method, string(content))
+		}
+	}
+}
+
+func (c *decoratedCollection) logDurationSimple(method string, startTime time.Duration, err error) {
+	logDuration(c.name, method, startTime, err)
+}
+
+func (p keepablePromise) accept(err error) error {
+	p.promise.Accept()
+	p.log(err)
+	return err
+}
+
+func (p keepablePromise) keep(err error) error {
+	if acceptable(err) {
+		p.promise.Accept()
+	} else {
+		p.promise.Reject(err.Error())
+	}
+
+	p.log(err)
+	return err
+}
+
+func acceptable(err error) bool {
+	return err == nil || err == mongo.ErrNoDocuments || err == mongo.ErrNilValue ||
+		err == mongo.ErrNilDocument || err == mongo.ErrNilCursor || err == mongo.ErrEmptySlice ||
+		// session errors
+		err == session.ErrSessionEnded || err == session.ErrNoTransactStarted ||
+		err == session.ErrTransactInProgress || err == session.ErrAbortAfterCommit ||
+		err == session.ErrAbortTwice || err == session.ErrCommitAfterAbort ||
+		err == session.ErrUnackWCUnsupported || err == session.ErrSnapshotTransaction
+}
+
+func startSpan(ctx context.Context) (context.Context, tracesdk.Span) {
+	tracer := otel.GetTracerProvider().Tracer(trace.TraceName)
+	return tracer.Start(ctx, spanName)
+}

core/stores/mon/collection_test.go

@@ -0,0 +1,608 @@
+package mon
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/zeromicro/go-zero/core/breaker"
+	"github.com/zeromicro/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/core/stringx"
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/bson/primitive"
+	"go.mongodb.org/mongo-driver/mongo"
+	"go.mongodb.org/mongo-driver/mongo/integration/mtest"
+	mopt "go.mongodb.org/mongo-driver/mongo/options"
+)
+
+var errDummy = errors.New("dummy")
+
+func init() {
+	logx.Disable()
+}
+
+func TestKeepPromise_accept(t *testing.T) {
+	p := new(mockPromise)
+	kp := keepablePromise{
+		promise: p,
+		log:     func(error) {},
+	}
+	assert.Nil(t, kp.accept(nil))
+	assert.Equal(t, ErrNotFound, kp.accept(ErrNotFound))
+}
+
+func TestKeepPromise_keep(t *testing.T) {
+	tests := []struct {
+		err      error
+		accepted bool
+		reason   string
+	}{
+		{
+			err:      nil,
+			accepted: true,
+			reason:   "",
+		},
+		{
+			err:      ErrNotFound,
+			accepted: true,
+			reason:   "",
+		},
+		{
+			err:      errors.New("any"),
+			accepted: false,
+			reason:   "any",
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(stringx.RandId(), func(t *testing.T) {
+			p := new(mockPromise)
+			kp := keepablePromise{
+				promise: p,
+				log:     func(error) {},
+			}
+			assert.Equal(t, test.err, kp.keep(test.err))
+			assert.Equal(t, test.accepted, p.accepted)
+			assert.Equal(t, test.reason, p.reason)
+		})
+	}
+}
+
+func TestNewCollection(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+	mt.Run("test", func(mt *mtest.T) {
+		coll := mt.Coll
+		assert.NotNil(t, coll)
+		col := newCollection(coll, breaker.GetBreaker("localhost"))
+		assert.Equal(t, t.Name()+"/test", col.(*decoratedCollection).name)
+	})
+}
+
+func TestCollection_Aggregate(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+	mt.Run("test", func(mt *mtest.T) {
+		coll := mt.Coll
+		assert.NotNil(t, coll)
+		col := newCollection(coll, breaker.GetBreaker("localhost"))
+		ns := mt.Coll.Database().Name() + "." + mt.Coll.Name()
+		aggRes := mtest.CreateCursorResponse(1, ns, mtest.FirstBatch)
+		mt.AddMockResponses(aggRes)
+		assert.Equal(t, t.Name()+"/test", col.(*decoratedCollection).name)
+		cursor, err := col.Aggregate(context.Background(), mongo.Pipeline{}, mopt.Aggregate())
+		assert.Nil(t, err)
+		cursor.Close(context.Background())
+	})
+}
+
+func TestCollection_BulkWrite(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		c := decoratedCollection{
+			Collection: mt.Coll,
+			brk:        breaker.NewBreaker(),
+		}
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "ok", Value: 1}}...))
+		res, err := c.BulkWrite(context.Background(), []mongo.WriteModel{
+			mongo.NewInsertOneModel().SetDocument(bson.D{{Key: "foo", Value: 1}})},
+		)
+		assert.Nil(t, err)
+		assert.NotNil(t, res)
+		c.brk = new(dropBreaker)
+		_, err = c.BulkWrite(context.Background(), []mongo.WriteModel{
+			mongo.NewInsertOneModel().SetDocument(bson.D{{Key: "foo", Value: 1}})},
+		)
+		assert.Equal(t, errDummy, err)
+	})
+}
+
+func TestCollection_CountDocuments(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		c := decoratedCollection{
+			Collection: mt.Coll,
+			brk:        breaker.NewBreaker(),
+		}
+		mt.AddMockResponses(mtest.CreateCursorResponse(
+			1,
+			"DBName.CollectionName",
+			mtest.FirstBatch,
+			bson.D{
+				{Key: "n", Value: 1},
+			}))
+		res, err := c.CountDocuments(context.Background(), bson.D{})
+		assert.Nil(t, err)
+		assert.Equal(t, int64(1), res)
+
+		c.brk = new(dropBreaker)
+		_, err = c.CountDocuments(context.Background(), bson.D{{Key: "foo", Value: 1}})
+		assert.Equal(t, errDummy, err)
+	})
+}
+
+func TestDecoratedCollection_DeleteMany(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		c := decoratedCollection{
+			Collection: mt.Coll,
+			brk:        breaker.NewBreaker(),
+		}
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
+		res, err := c.DeleteMany(context.Background(), bson.D{})
+		assert.Nil(t, err)
+		assert.Equal(t, int64(1), res.DeletedCount)
+
+		c.brk = new(dropBreaker)
+		_, err = c.DeleteMany(context.Background(), bson.D{{Key: "foo", Value: 1}})
+		assert.Equal(t, errDummy, err)
+	})
+}
+
+func TestCollection_Distinct(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		c := decoratedCollection{
+			Collection: mt.Coll,
+			brk:        breaker.NewBreaker(),
+		}
+		mt.AddMockResponses(bson.D{{Key: "ok", Value: 1}, {Key: "values", Value: []int{1}}})
+		resp, err := c.Distinct(context.Background(), "foo", bson.D{})
+		assert.Nil(t, err)
+		assert.Equal(t, 1, len(resp))
+
+		c.brk = new(dropBreaker)
+		_, err = c.Distinct(context.Background(), "foo", bson.D{{Key: "foo", Value: 1}})
+		assert.Equal(t, errDummy, err)
+	})
+}
+
+func TestCollection_EstimatedDocumentCount(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		c := decoratedCollection{
+			Collection: mt.Coll,
+			brk:        breaker.NewBreaker(),
+		}
+		mt.AddMockResponses(bson.D{{Key: "ok", Value: 1}, {Key: "n", Value: 1}})
+		res, err := c.EstimatedDocumentCount(context.Background())
+		assert.Nil(t, err)
+		assert.Equal(t, int64(1), res)
+
+		c.brk = new(dropBreaker)
+		_, err = c.EstimatedDocumentCount(context.Background())
+		assert.Equal(t, errDummy, err)
+	})
+}
+
+func TestCollectionFind(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		c := decoratedCollection{
+			Collection: mt.Coll,
+			brk:        breaker.NewBreaker(),
+		}
+		find := mtest.CreateCursorResponse(
+			1,
+			"DBName.CollectionName",
+			mtest.FirstBatch,
+			bson.D{
+				{Key: "name", Value: "John"},
+			})
+		getMore := mtest.CreateCursorResponse(
+			1,
+			"DBName.CollectionName",
+			mtest.NextBatch,
+			bson.D{
+				{Key: "name", Value: "Mary"},
+			})
+		killCursors := mtest.CreateCursorResponse(
+			0,
+			"DBName.CollectionName",
+			mtest.NextBatch)
+		mt.AddMockResponses(find, getMore, killCursors)
+		filter := bson.D{{Key: "x", Value: 1}}
+		cursor, err := c.Find(context.Background(), filter, mopt.Find())
+		assert.Nil(t, err)
+		defer cursor.Close(context.Background())
+
+		var val []struct {
+			ID   primitive.ObjectID `bson:"_id"`
+			Name string             `bson:"name"`
+		}
+		assert.Nil(t, cursor.All(context.Background(), &val))
+		assert.Equal(t, 2, len(val))
+		assert.Equal(t, "John", val[0].Name)
+		assert.Equal(t, "Mary", val[1].Name)
+
+		c.brk = new(dropBreaker)
+		_, err = c.Find(context.Background(), filter, mopt.Find())
+		assert.Equal(t, errDummy, err)
+	})
+}
+
+func TestCollectionFindOne(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		c := decoratedCollection{
+			Collection: mt.Coll,
+			brk:        breaker.NewBreaker(),
+		}
+		find := mtest.CreateCursorResponse(
+			1,
+			"DBName.CollectionName",
+			mtest.FirstBatch,
+			bson.D{
+				{Key: "name", Value: "John"},
+			})
+		getMore := mtest.CreateCursorResponse(
+			1,
+			"DBName.CollectionName",
+			mtest.NextBatch,
+			bson.D{
+				{Key: "name", Value: "Mary"},
+			})
+		killCursors := mtest.CreateCursorResponse(
+			0,
+			"DBName.CollectionName",
+			mtest.NextBatch)
+		mt.AddMockResponses(find, getMore, killCursors)
+		filter := bson.D{{Key: "x", Value: 1}}
+		resp, err := c.FindOne(context.Background(), filter)
+		assert.Nil(t, err)
+		var val struct {
+			ID   primitive.ObjectID `bson:"_id"`
+			Name string             `bson:"name"`
+		}
+		assert.Nil(t, resp.Decode(&val))
+		assert.Equal(t, "John", val.Name)
+
+		c.brk = new(dropBreaker)
+		_, err = c.FindOne(context.Background(), filter)
+		assert.Equal(t, errDummy, err)
+	})
+}
+
+func TestCollection_FindOneAndDelete(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		c := decoratedCollection{
+			Collection: mt.Coll,
+			brk:        breaker.NewBreaker(),
+		}
+		filter := bson.D{}
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{}...))
+		_, err := c.FindOneAndDelete(context.Background(), filter, mopt.FindOneAndDelete())
+		assert.Equal(t, mongo.ErrNoDocuments, err)
+
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
+			{Key: "value", Value: bson.D{{Key: "name", Value: "John"}}},
+		}...))
+		resp, err := c.FindOneAndDelete(context.Background(), filter, mopt.FindOneAndDelete())
+		assert.Nil(t, err)
+		var val struct {
+			Name string `bson:"name"`
+		}
+		assert.Nil(t, resp.Decode(&val))
+		assert.Equal(t, "John", val.Name)
+
+		c.brk = new(dropBreaker)
+		_, err = c.FindOneAndDelete(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
+		assert.Equal(t, errDummy, err)
+	})
+}
+
+func TestCollection_FindOneAndReplace(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		c := decoratedCollection{
+			Collection: mt.Coll,
+			brk:        breaker.NewBreaker(),
+		}
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{}...))
+		filter := bson.D{{Key: "x", Value: 1}}
+		replacement := bson.D{{Key: "x", Value: 2}}
+		opts := mopt.FindOneAndReplace().SetUpsert(true)
+		_, err := c.FindOneAndReplace(context.Background(), filter, replacement, opts)
+		assert.Equal(t, mongo.ErrNoDocuments, err)
+		mt.AddMockResponses(bson.D{{Key: "ok", Value: 1}, {Key: "value", Value: bson.D{
+			{Key: "name", Value: "John"},
+		}}})
+		resp, err := c.FindOneAndReplace(context.Background(), filter, replacement, opts)
+		assert.Nil(t, err)
+		var val struct {
+			Name string `bson:"name"`
+		}
+		assert.Nil(t, resp.Decode(&val))
+		assert.Equal(t, "John", val.Name)
+
+		c.brk = new(dropBreaker)
+		_, err = c.FindOneAndReplace(context.Background(), filter, replacement, opts)
+		assert.Equal(t, errDummy, err)
+	})
+}
+
+func TestCollection_FindOneAndUpdate(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		c := decoratedCollection{
+			Collection: mt.Coll,
+			brk:        breaker.NewBreaker(),
+		}
+		mt.AddMockResponses(bson.D{{Key: "ok", Value: 1}})
+		filter := bson.D{{Key: "x", Value: 1}}
+		update := bson.D{{Key: "$x", Value: 2}}
+		opts := mopt.FindOneAndUpdate().SetUpsert(true)
+		_, err := c.FindOneAndUpdate(context.Background(), filter, update, opts)
+		assert.Equal(t, mongo.ErrNoDocuments, err)
+
+		mt.AddMockResponses(bson.D{{Key: "ok", Value: 1}, {Key: "value", Value: bson.D{
+			{Key: "name", Value: "John"},
+		}}})
+		resp, err := c.FindOneAndUpdate(context.Background(), filter, update, opts)
+		assert.Nil(t, err)
+		var val struct {
+			Name string `bson:"name"`
+		}
+		assert.Nil(t, resp.Decode(&val))
+		assert.Equal(t, "John", val.Name)
+
+		c.brk = new(dropBreaker)
+		_, err = c.FindOneAndUpdate(context.Background(), filter, update, opts)
+		assert.Equal(t, errDummy, err)
+	})
+}
+
+func TestCollection_InsertOne(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		c := decoratedCollection{
+			Collection: mt.Coll,
+			brk:        breaker.NewBreaker(),
+		}
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "ok", Value: 1}}...))
+		res, err := c.InsertOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
+		assert.Nil(t, err)
+		assert.NotNil(t, res)
+
+		c.brk = new(dropBreaker)
+		_, err = c.InsertOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
+		assert.Equal(t, errDummy, err)
+	})
+}
+
+func TestCollection_InsertMany(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		c := decoratedCollection{
+			Collection: mt.Coll,
+			brk:        breaker.NewBreaker(),
+		}
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "ok", Value: 1}}...))
+		res, err := c.InsertMany(context.Background(), []interface{}{
+			bson.D{{Key: "foo", Value: "bar"}},
+			bson.D{{Key: "foo", Value: "baz"}},
+		})
+		assert.Nil(t, err)
+		assert.NotNil(t, res)
+		assert.Equal(t, 2, len(res.InsertedIDs))
+
+		c.brk = new(dropBreaker)
+		_, err = c.InsertMany(context.Background(), []interface{}{bson.D{{Key: "foo", Value: "bar"}}})
+		assert.Equal(t, errDummy, err)
+	})
+}
+
+func TestCollection_Remove(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		c := decoratedCollection{
+			Collection: mt.Coll,
+			brk:        breaker.NewBreaker(),
+		}
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
+		res, err := c.DeleteOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
+		assert.Nil(t, err)
+		assert.Equal(t, int64(1), res.DeletedCount)
+
+		c.brk = new(dropBreaker)
+		_, err = c.DeleteOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
+		assert.Equal(t, errDummy, err)
+	})
+}
+
+func TestCollectionRemoveAll(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		c := decoratedCollection{
+			Collection: mt.Coll,
+			brk:        breaker.NewBreaker(),
+		}
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
+		res, err := c.DeleteMany(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
+		assert.Nil(t, err)
+		assert.Equal(t, int64(1), res.DeletedCount)
+
+		c.brk = new(dropBreaker)
+		_, err = c.DeleteMany(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
+		assert.Equal(t, errDummy, err)
+	})
+}
+
+func TestCollection_ReplaceOne(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		c := decoratedCollection{
+			Collection: mt.Coll,
+			brk:        breaker.NewBreaker(),
+		}
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
+		res, err := c.ReplaceOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}},
+			bson.D{{Key: "foo", Value: "baz"}},
+		)
+		assert.Nil(t, err)
+		assert.Equal(t, int64(1), res.MatchedCount)
+
+		c.brk = new(dropBreaker)
+		_, err = c.ReplaceOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}},
+			bson.D{{Key: "foo", Value: "baz"}})
+		assert.Equal(t, errDummy, err)
+	})
+}
+
+func TestCollection_UpdateOne(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		c := decoratedCollection{
+			Collection: mt.Coll,
+			brk:        breaker.NewBreaker(),
+		}
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
+		resp, err := c.UpdateOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}},
+			bson.D{{Key: "$set", Value: bson.D{{Key: "baz", Value: "qux"}}}})
+		assert.Nil(t, err)
+		assert.Equal(t, int64(1), resp.MatchedCount)
+
+		c.brk = new(dropBreaker)
+		_, err = c.UpdateOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}},
+			bson.D{{Key: "$set", Value: bson.D{{Key: "baz", Value: "qux"}}}})
+		assert.Equal(t, errDummy, err)
+	})
+}
+
+func TestCollection_UpdateByID(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		c := decoratedCollection{
+			Collection: mt.Coll,
+			brk:        breaker.NewBreaker(),
+		}
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
+		resp, err := c.UpdateByID(context.Background(), primitive.NewObjectID(),
+			bson.D{{Key: "$set", Value: bson.D{{Key: "baz", Value: "qux"}}}})
+		assert.Nil(t, err)
+		assert.Equal(t, int64(1), resp.MatchedCount)
+
+		c.brk = new(dropBreaker)
+		_, err = c.UpdateByID(context.Background(), primitive.NewObjectID(),
+			bson.D{{Key: "$set", Value: bson.D{{Key: "baz", Value: "qux"}}}})
+		assert.Equal(t, errDummy, err)
+	})
+}
+
+func TestCollection_UpdateMany(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		c := decoratedCollection{
+			Collection: mt.Coll,
+			brk:        breaker.NewBreaker(),
+		}
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
+		resp, err := c.UpdateMany(context.Background(), bson.D{{Key: "foo", Value: "bar"}},
+			bson.D{{Key: "$set", Value: bson.D{{Key: "baz", Value: "qux"}}}})
+		assert.Nil(t, err)
+		assert.Equal(t, int64(1), resp.MatchedCount)
+
+		c.brk = new(dropBreaker)
+		_, err = c.UpdateMany(context.Background(), bson.D{{Key: "foo", Value: "bar"}},
+			bson.D{{Key: "$set", Value: bson.D{{Key: "baz", Value: "qux"}}}})
+		assert.Equal(t, errDummy, err)
+	})
+}
+
+type mockPromise struct {
+	accepted bool
+	reason   string
+}
+
+func (p *mockPromise) Accept() {
+	p.accepted = true
+}
+
+func (p *mockPromise) Reject(reason string) {
+	p.reason = reason
+}
+
+type dropBreaker struct{}
+
+func (d *dropBreaker) Name() string {
+	return "dummy"
+}
+
+func (d *dropBreaker) Allow() (breaker.Promise, error) {
+	return nil, errDummy
+}
+
+func (d *dropBreaker) Do(req func() error) error {
+	return nil
+}
+
+func (d *dropBreaker) DoWithAcceptable(req func() error, acceptable breaker.Acceptable) error {
+	return errDummy
+}
+
+func (d *dropBreaker) DoWithFallback(req func() error, fallback func(err error) error) error {
+	return nil
+}
+
+func (d *dropBreaker) DoWithFallbackAcceptable(_ func() error, _ func(err error) error,
+	_ breaker.Acceptable) error {
+	return nil
+}

core/stores/mon/model.go

@@ -0,0 +1,214 @@
+package mon
+
+import (
+	"context"
+	"log"
+	"strings"
+
+	"github.com/zeromicro/go-zero/core/breaker"
+	"github.com/zeromicro/go-zero/core/timex"
+	"go.mongodb.org/mongo-driver/mongo"
+	mopt "go.mongodb.org/mongo-driver/mongo/options"
+)
+
+type (
+	// Model is a mongodb store model that represents a collection.
+	Model struct {
+		Collection
+		name string
+		cli  *mongo.Client
+		brk  breaker.Breaker
+		opts []Option
+	}
+
+	wrappedSession struct {
+		mongo.Session
+		brk breaker.Breaker
+	}
+)
+
+// MustNewModel returns a Model, exits on errors.
+func MustNewModel(uri, db, collection string, opts ...Option) *Model {
+	model, err := NewModel(uri, db, collection, opts...)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	return model
+}
+
+// NewModel returns a Model.
+func NewModel(uri, db, collection string, opts ...Option) (*Model, error) {
+	cli, err := getClient(uri)
+	if err != nil {
+		return nil, err
+	}
+
+	name := strings.Join([]string{uri, collection}, "/")
+	brk := breaker.GetBreaker(uri)
+	coll := newCollection(cli.Database(db).Collection(collection), brk)
+	return newModel(name, cli, coll, brk, opts...), nil
+}
+
+func newModel(name string, cli *mongo.Client, coll Collection, brk breaker.Breaker,
+	opts ...Option) *Model {
+	return &Model{
+		name:       name,
+		Collection: coll,
+		cli:        cli,
+		brk:        brk,
+		opts:       opts,
+	}
+}
+
+// StartSession starts a new session.
+func (m *Model) StartSession(opts ...*mopt.SessionOptions) (sess mongo.Session, err error) {
+	err = m.brk.DoWithAcceptable(func() error {
+		starTime := timex.Now()
+		defer func() {
+			logDuration(m.name, "StartSession", starTime, err)
+		}()
+
+		session, sessionErr := m.cli.StartSession(opts...)
+		if sessionErr != nil {
+			return sessionErr
+		}
+
+		sess = &wrappedSession{
+			Session: session,
+			brk:     m.brk,
+		}
+
+		return nil
+	}, acceptable)
+	return
+}
+
+// Aggregate executes an aggregation pipeline.
+func (m *Model) Aggregate(ctx context.Context, v, pipeline interface{}, opts ...*mopt.AggregateOptions) error {
+	cur, err := m.Collection.Aggregate(ctx, pipeline, opts...)
+	if err != nil {
+		return err
+	}
+	defer cur.Close(ctx)
+
+	return cur.All(ctx, v)
+}
+
+// DeleteMany deletes documents that match the filter.
+func (m *Model) DeleteMany(ctx context.Context, filter interface{}, opts ...*mopt.DeleteOptions) (int64, error) {
+	res, err := m.Collection.DeleteMany(ctx, filter, opts...)
+	if err != nil {
+		return 0, err
+	}
+
+	return res.DeletedCount, nil
+}
+
+// DeleteOne deletes the first document that matches the filter.
+func (m *Model) DeleteOne(ctx context.Context, filter interface{}, opts ...*mopt.DeleteOptions) (int64, error) {
+	res, err := m.Collection.DeleteOne(ctx, filter, opts...)
+	if err != nil {
+		return 0, err
+	}
+
+	return res.DeletedCount, nil
+}
+
+// Find finds documents that match the filter.
+func (m *Model) Find(ctx context.Context, v, filter interface{}, opts ...*mopt.FindOptions) error {
+	cur, err := m.Collection.Find(ctx, filter, opts...)
+	if err != nil {
+		return err
+	}
+	defer cur.Close(ctx)
+
+	return cur.All(ctx, v)
+}
+
+// FindOne finds the first document that matches the filter.
+func (m *Model) FindOne(ctx context.Context, v, filter interface{}, opts ...*mopt.FindOneOptions) error {
+	res, err := m.Collection.FindOne(ctx, filter, opts...)
+	if err != nil {
+		return err
+	}
+
+	return res.Decode(v)
+}
+
+// FindOneAndDelete finds a single document and deletes it.
+func (m *Model) FindOneAndDelete(ctx context.Context, v, filter interface{},
+	opts ...*mopt.FindOneAndDeleteOptions) error {
+	res, err := m.Collection.FindOneAndDelete(ctx, filter, opts...)
+	if err != nil {
+		return err
+	}
+
+	return res.Decode(v)
+}
+
+// FindOneAndReplace finds a single document and replaces it.
+func (m *Model) FindOneAndReplace(ctx context.Context, v, filter interface{}, replacement interface{},
+	opts ...*mopt.FindOneAndReplaceOptions) error {
+	res, err := m.Collection.FindOneAndReplace(ctx, filter, replacement, opts...)
+	if err != nil {
+		return err
+	}
+
+	return res.Decode(v)
+}
+
+// FindOneAndUpdate finds a single document and updates it.
+func (m *Model) FindOneAndUpdate(ctx context.Context, v, filter interface{}, update interface{},
+	opts ...*mopt.FindOneAndUpdateOptions) error {
+	res, err := m.Collection.FindOneAndUpdate(ctx, filter, update, opts...)
+	if err != nil {
+		return err
+	}
+
+	return res.Decode(v)
+}
+
+func (w *wrappedSession) AbortTransaction(ctx context.Context) error {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
+	return w.brk.DoWithAcceptable(func() error {
+		return w.Session.AbortTransaction(ctx)
+	}, acceptable)
+}
+
+func (w *wrappedSession) CommitTransaction(ctx context.Context) error {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
+	return w.brk.DoWithAcceptable(func() error {
+		return w.Session.CommitTransaction(ctx)
+	}, acceptable)
+}
+
+func (w *wrappedSession) WithTransaction(
+	ctx context.Context,
+	fn func(sessCtx mongo.SessionContext) (interface{}, error),
+	opts ...*mopt.TransactionOptions,
+) (res interface{}, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
+	err = w.brk.DoWithAcceptable(func() error {
+		res, err = w.Session.WithTransaction(ctx, fn, opts...)
+		return err
+	}, acceptable)
+
+	return
+}
+
+func (w *wrappedSession) EndSession(ctx context.Context) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
+	_ = w.brk.DoWithAcceptable(func() error {
+		w.Session.EndSession(ctx)
+		return nil
+	}, acceptable)
+}

core/stores/mon/model_test.go

@@ -0,0 +1,243 @@
+package mon
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/mongo"
+	"go.mongodb.org/mongo-driver/mongo/integration/mtest"
+)
+
+func TestModel_StartSession(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		m := createModel(mt)
+		sess, err := m.StartSession()
+		assert.Nil(t, err)
+		defer sess.EndSession(context.Background())
+
+		_, err = sess.WithTransaction(context.Background(), func(sessCtx mongo.SessionContext) (interface{}, error) {
+			_ = sessCtx.StartTransaction()
+			sessCtx.Client().Database("1")
+			sessCtx.EndSession(context.Background())
+			return nil, nil
+		})
+		assert.Nil(t, err)
+		assert.NoError(t, sess.CommitTransaction(context.Background()))
+		assert.Error(t, sess.AbortTransaction(context.Background()))
+	})
+}
+
+func TestModel_Aggregate(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		m := createModel(mt)
+		find := mtest.CreateCursorResponse(
+			1,
+			"DBName.CollectionName",
+			mtest.FirstBatch,
+			bson.D{
+				{Key: "name", Value: "John"},
+			})
+		getMore := mtest.CreateCursorResponse(
+			1,
+			"DBName.CollectionName",
+			mtest.NextBatch,
+			bson.D{
+				{Key: "name", Value: "Mary"},
+			})
+		killCursors := mtest.CreateCursorResponse(
+			0,
+			"DBName.CollectionName",
+			mtest.NextBatch)
+		mt.AddMockResponses(find, getMore, killCursors)
+		var result []interface{}
+		err := m.Aggregate(context.Background(), &result, mongo.Pipeline{})
+		assert.Nil(t, err)
+		assert.Equal(t, 2, len(result))
+		assert.Equal(t, "John", result[0].(bson.D).Map()["name"])
+		assert.Equal(t, "Mary", result[1].(bson.D).Map()["name"])
+
+		triggerBreaker(m)
+		assert.Equal(t, errDummy, m.Aggregate(context.Background(), &result, mongo.Pipeline{}))
+	})
+}
+
+func TestModel_DeleteMany(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		m := createModel(mt)
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
+		val, err := m.DeleteMany(context.Background(), bson.D{})
+		assert.Nil(t, err)
+		assert.Equal(t, int64(1), val)
+
+		triggerBreaker(m)
+		_, err = m.DeleteMany(context.Background(), bson.D{})
+		assert.Equal(t, errDummy, err)
+	})
+}
+
+func TestModel_DeleteOne(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		m := createModel(mt)
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
+		val, err := m.DeleteOne(context.Background(), bson.D{})
+		assert.Nil(t, err)
+		assert.Equal(t, int64(1), val)
+
+		triggerBreaker(m)
+		_, err = m.DeleteOne(context.Background(), bson.D{})
+		assert.Equal(t, errDummy, err)
+	})
+}
+
+func TestModel_Find(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		m := createModel(mt)
+		find := mtest.CreateCursorResponse(
+			1,
+			"DBName.CollectionName",
+			mtest.FirstBatch,
+			bson.D{
+				{Key: "name", Value: "John"},
+			})
+		getMore := mtest.CreateCursorResponse(
+			1,
+			"DBName.CollectionName",
+			mtest.NextBatch,
+			bson.D{
+				{Key: "name", Value: "Mary"},
+			})
+		killCursors := mtest.CreateCursorResponse(
+			0,
+			"DBName.CollectionName",
+			mtest.NextBatch)
+		mt.AddMockResponses(find, getMore, killCursors)
+		var result []interface{}
+		err := m.Find(context.Background(), &result, bson.D{})
+		assert.Nil(t, err)
+		assert.Equal(t, 2, len(result))
+		assert.Equal(t, "John", result[0].(bson.D).Map()["name"])
+		assert.Equal(t, "Mary", result[1].(bson.D).Map()["name"])
+
+		triggerBreaker(m)
+		assert.Equal(t, errDummy, m.Find(context.Background(), &result, bson.D{}))
+	})
+}
+
+func TestModel_FindOne(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		m := createModel(mt)
+		find := mtest.CreateCursorResponse(
+			1,
+			"DBName.CollectionName",
+			mtest.FirstBatch,
+			bson.D{
+				{Key: "name", Value: "John"},
+			})
+		killCursors := mtest.CreateCursorResponse(
+			0,
+			"DBName.CollectionName",
+			mtest.NextBatch)
+		mt.AddMockResponses(find, killCursors)
+		var result bson.D
+		err := m.FindOne(context.Background(), &result, bson.D{})
+		assert.Nil(t, err)
+		assert.Equal(t, "John", result.Map()["name"])
+
+		triggerBreaker(m)
+		assert.Equal(t, errDummy, m.FindOne(context.Background(), &result, bson.D{}))
+	})
+}
+
+func TestModel_FindOneAndDelete(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		m := createModel(mt)
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
+			{Key: "value", Value: bson.D{{Key: "name", Value: "John"}}},
+		}...))
+		var result bson.D
+		err := m.FindOneAndDelete(context.Background(), &result, bson.D{})
+		assert.Nil(t, err)
+		assert.Equal(t, "John", result.Map()["name"])
+
+		triggerBreaker(m)
+		assert.Equal(t, errDummy, m.FindOneAndDelete(context.Background(), &result, bson.D{}))
+	})
+}
+
+func TestModel_FindOneAndReplace(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		m := createModel(mt)
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
+			{Key: "value", Value: bson.D{{Key: "name", Value: "John"}}},
+		}...))
+		var result bson.D
+		err := m.FindOneAndReplace(context.Background(), &result, bson.D{}, bson.D{
+			{Key: "name", Value: "Mary"},
+		})
+		assert.Nil(t, err)
+		assert.Equal(t, "John", result.Map()["name"])
+
+		triggerBreaker(m)
+		assert.Equal(t, errDummy, m.FindOneAndReplace(context.Background(), &result, bson.D{}, bson.D{
+			{Key: "name", Value: "Mary"},
+		}))
+	})
+}
+
+func TestModel_FindOneAndUpdate(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		m := createModel(mt)
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
+			{Key: "value", Value: bson.D{{Key: "name", Value: "John"}}},
+		}...))
+		var result bson.D
+		err := m.FindOneAndUpdate(context.Background(), &result, bson.D{}, bson.D{
+			{Key: "$set", Value: bson.D{{Key: "name", Value: "Mary"}}},
+		})
+		assert.Nil(t, err)
+		assert.Equal(t, "John", result.Map()["name"])
+
+		triggerBreaker(m)
+		assert.Equal(t, errDummy, m.FindOneAndUpdate(context.Background(), &result, bson.D{}, bson.D{
+			{Key: "$set", Value: bson.D{{Key: "name", Value: "Mary"}}},
+		}))
+	})
+}
+
+func createModel(mt *mtest.T) *Model {
+	Inject(mt.Name(), mt.Client)
+	return MustNewModel(mt.Name(), mt.DB.Name(), mt.Coll.Name())
+}
+
+func triggerBreaker(m *Model) {
+	m.Collection.(*decoratedCollection).brk = new(dropBreaker)
+}

core/stores/mon/options.go

@@ -0,0 +1,29 @@
+package mon
+
+import (
+	"time"
+
+	"github.com/zeromicro/go-zero/core/syncx"
+)
+
+var slowThreshold = syncx.ForAtomicDuration(defaultSlowThreshold)
+
+type (
+	options struct {
+		timeout time.Duration
+	}
+
+	// Option defines the method to customize a mongo model.
+	Option func(opts *options)
+)
+
+// SetSlowThreshold sets the slow threshold.
+func SetSlowThreshold(threshold time.Duration) {
+	slowThreshold.Set(threshold)
+}
+
+func defaultOptions() *options {
+	return &options{
+		timeout: defaultTimeout,
+	}
+}

core/stores/mon/options_test.go

@@ -0,0 +1,18 @@
+package mon
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestSetSlowThreshold(t *testing.T) {
+	assert.Equal(t, defaultSlowThreshold, slowThreshold.Load())
+	SetSlowThreshold(time.Second)
+	assert.Equal(t, time.Second, slowThreshold.Load())
+}
+
+func TestDefaultOptions(t *testing.T) {
+	assert.Equal(t, defaultTimeout, defaultOptions().timeout)
+}

core/stores/mon/util.go

@@ -0,0 +1,25 @@
+package mon
+
+import (
+	"strings"
+	"time"
+
+	"github.com/zeromicro/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/core/timex"
+)
+
+const mongoAddrSep = ","
+
+// FormatAddr formats mongo hosts to a string.
+func FormatAddr(hosts []string) string {
+	return strings.Join(hosts, mongoAddrSep)
+}
+
+func logDuration(name, method string, startTime time.Duration, err error) {
+	duration := timex.Since(startTime)
+	if err != nil {
+		logx.WithDuration(duration).Infof("mongo(%s) - %s - fail(%s)", name, method, err.Error())
+	} else {
+		logx.WithDuration(duration).Infof("mongo(%s) - %s - ok", name, method)
+	}
+}

core/stores/mon/util_test.go

@@ -0,0 +1,35 @@
+package mon
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestFormatAddrs(t *testing.T) {
+	tests := []struct {
+		addrs  []string
+		expect string
+	}{
+		{
+			addrs:  []string{"a", "b"},
+			expect: "a,b",
+		},
+		{
+			addrs:  []string{"a", "b", "c"},
+			expect: "a,b,c",
+		},
+		{
+			addrs:  []string{},
+			expect: "",
+		},
+		{
+			addrs:  nil,
+			expect: "",
+		},
+	}
+
+	for _, test := range tests {
+		assert.Equal(t, test.expect, FormatAddr(test.addrs))
+	}
+}

core/stores/monc/cachedmodel.go

@@ -0,0 +1,281 @@
+package monc
+
+import (
+	"context"
+	"log"
+
+	"github.com/zeromicro/go-zero/core/stores/cache"
+	"github.com/zeromicro/go-zero/core/stores/mon"
+	"github.com/zeromicro/go-zero/core/stores/redis"
+	"github.com/zeromicro/go-zero/core/syncx"
+	"go.mongodb.org/mongo-driver/mongo"
+	mopt "go.mongodb.org/mongo-driver/mongo/options"
+)
+
+var (
+	// ErrNotFound is an alias of mongo.ErrNoDocuments.
+	ErrNotFound = mongo.ErrNoDocuments
+
+	// can't use one SingleFlight per conn, because multiple conns may share the same cache key.
+	singleFlight = syncx.NewSingleFlight()
+	stats        = cache.NewStat("monc")
+)
+
+// A Model is a mongo model that built with cache capability.
+type Model struct {
+	*mon.Model
+	cache cache.Cache
+}
+
+// MustNewModel returns a Model with a cache cluster, exists on errors.
+func MustNewModel(uri, db, collection string, c cache.CacheConf, opts ...cache.Option) *Model {
+	model, err := NewModel(uri, db, collection, c, opts...)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	return model
+}
+
+// MustNewNodeModel returns a Model with a cache node, exists on errors.
+func MustNewNodeModel(uri, db, collection string, rds *redis.Redis, opts ...cache.Option) *Model {
+	model, err := NewNodeModel(uri, db, collection, rds, opts...)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	return model
+}
+
+// NewModel returns a Model with a cache cluster.
+func NewModel(uri, db, collection string, conf cache.CacheConf, opts ...cache.Option) (*Model, error) {
+	c := cache.New(conf, singleFlight, stats, mongo.ErrNoDocuments, opts...)
+	return NewModelWithCache(uri, db, collection, c)
+}
+
+// NewModelWithCache returns a Model with a custom cache.
+func NewModelWithCache(uri, db, collection string, c cache.Cache) (*Model, error) {
+	return newModel(uri, db, collection, c)
+}
+
+// NewNodeModel returns a Model with a cache node.
+func NewNodeModel(uri, db, collection string, rds *redis.Redis, opts ...cache.Option) (*Model, error) {
+	c := cache.NewNode(rds, singleFlight, stats, mongo.ErrNoDocuments, opts...)
+	return NewModelWithCache(uri, db, collection, c)
+}
+
+// newModel returns a Model with the given cache.
+func newModel(uri, db, collection string, c cache.Cache) (*Model, error) {
+	model, err := mon.NewModel(uri, db, collection)
+	if err != nil {
+		return nil, err
+	}
+
+	return &Model{
+		Model: model,
+		cache: c,
+	}, nil
+}
+
+// DelCache deletes the cache with given keys.
+func (mm *Model) DelCache(ctx context.Context, keys ...string) error {
+	return mm.cache.DelCtx(ctx, keys...)
+}
+
+// DeleteOne deletes the document with given filter, and remove it from cache.
+func (mm *Model) DeleteOne(ctx context.Context, key string, filter interface{},
+	opts ...*mopt.DeleteOptions) (int64, error) {
+	val, err := mm.Model.DeleteOne(ctx, filter, opts...)
+	if err != nil {
+		return 0, err
+	}
+
+	if err := mm.DelCache(ctx, key); err != nil {
+		return 0, err
+	}
+
+	return val, nil
+}
+
+// DeleteOneNoCache deletes the document with given filter.
+func (mm *Model) DeleteOneNoCache(ctx context.Context, filter interface{},
+	opts ...*mopt.DeleteOptions) (int64, error) {
+	return mm.Model.DeleteOne(ctx, filter, opts...)
+}
+
+// FindOne unmarshals a record into v with given key and query.
+func (mm *Model) FindOne(ctx context.Context, key string, v, filter interface{},
+	opts ...*mopt.FindOneOptions) error {
+	return mm.cache.TakeCtx(ctx, v, key, func(v interface{}) error {
+		return mm.Model.FindOne(ctx, v, filter, opts...)
+	})
+}
+
+// FindOneNoCache unmarshals a record into v with query, without cache.
+func (mm *Model) FindOneNoCache(ctx context.Context, v, filter interface{},
+	opts ...*mopt.FindOneOptions) error {
+	return mm.Model.FindOne(ctx, v, filter, opts...)
+}
+
+// FindOneAndDelete deletes the document with given filter, and unmarshals it into v.
+func (mm *Model) FindOneAndDelete(ctx context.Context, key string, v, filter interface{},
+	opts ...*mopt.FindOneAndDeleteOptions) error {
+	if err := mm.Model.FindOneAndDelete(ctx, v, filter, opts...); err != nil {
+		return err
+	}
+
+	return mm.DelCache(ctx, key)
+}
+
+// FindOneAndDeleteNoCache deletes the document with given filter, and unmarshals it into v.
+func (mm *Model) FindOneAndDeleteNoCache(ctx context.Context, v, filter interface{},
+	opts ...*mopt.FindOneAndDeleteOptions) error {
+	return mm.Model.FindOneAndDelete(ctx, v, filter, opts...)
+}
+
+// FindOneAndReplace replaces the document with given filter with replacement, and unmarshals it into v.
+func (mm *Model) FindOneAndReplace(ctx context.Context, key string, v, filter interface{},
+	replacement interface{}, opts ...*mopt.FindOneAndReplaceOptions) error {
+	if err := mm.Model.FindOneAndReplace(ctx, v, filter, replacement, opts...); err != nil {
+		return err
+	}
+
+	return mm.DelCache(ctx, key)
+}
+
+// FindOneAndReplaceNoCache replaces the document with given filter with replacement, and unmarshals it into v.
+func (mm *Model) FindOneAndReplaceNoCache(ctx context.Context, v, filter interface{},
+	replacement interface{}, opts ...*mopt.FindOneAndReplaceOptions) error {
+	return mm.Model.FindOneAndReplace(ctx, v, filter, replacement, opts...)
+}
+
+// FindOneAndUpdate updates the document with given filter with update, and unmarshals it into v.
+func (mm *Model) FindOneAndUpdate(ctx context.Context, key string, v, filter interface{},
+	update interface{}, opts ...*mopt.FindOneAndUpdateOptions) error {
+	if err := mm.Model.FindOneAndUpdate(ctx, v, filter, update, opts...); err != nil {
+		return err
+	}
+
+	return mm.DelCache(ctx, key)
+}
+
+// FindOneAndUpdateNoCache updates the document with given filter with update, and unmarshals it into v.
+func (mm *Model) FindOneAndUpdateNoCache(ctx context.Context, v, filter interface{},
+	update interface{}, opts ...*mopt.FindOneAndUpdateOptions) error {
+	return mm.Model.FindOneAndUpdate(ctx, v, filter, update, opts...)
+}
+
+// GetCache unmarshal the cache into v with given key.
+func (mm *Model) GetCache(key string, v interface{}) error {
+	return mm.cache.Get(key, v)
+}
+
+// InsertOne inserts a single document into the collection, and remove the cache placeholder.
+func (mm *Model) InsertOne(ctx context.Context, key string, document interface{},
+	opts ...*mopt.InsertOneOptions) (*mongo.InsertOneResult, error) {
+	res, err := mm.Model.InsertOne(ctx, document, opts...)
+	if err != nil {
+		return nil, err
+	}
+
+	if err = mm.DelCache(ctx, key); err != nil {
+		return nil, err
+	}
+
+	return res, nil
+}
+
+// InsertOneNoCache inserts a single document into the collection.
+func (mm *Model) InsertOneNoCache(ctx context.Context, document interface{},
+	opts ...*mopt.InsertOneOptions) (*mongo.InsertOneResult, error) {
+	return mm.Model.InsertOne(ctx, document, opts...)
+}
+
+// ReplaceOne replaces a single document in the collection, and remove the cache.
+func (mm *Model) ReplaceOne(ctx context.Context, key string, filter interface{}, replacement interface{},
+	opts ...*mopt.ReplaceOptions) (*mongo.UpdateResult, error) {
+	res, err := mm.Model.ReplaceOne(ctx, filter, replacement, opts...)
+	if err != nil {
+		return nil, err
+	}
+
+	if err = mm.DelCache(ctx, key); err != nil {
+		return nil, err
+	}
+
+	return res, nil
+}
+
+// ReplaceOneNoCache replaces a single document in the collection.
+func (mm *Model) ReplaceOneNoCache(ctx context.Context, filter interface{}, replacement interface{},
+	opts ...*mopt.ReplaceOptions) (*mongo.UpdateResult, error) {
+	return mm.Model.ReplaceOne(ctx, filter, replacement, opts...)
+}
+
+// SetCache sets the cache with given key and value.
+func (mm *Model) SetCache(key string, v interface{}) error {
+	return mm.cache.Set(key, v)
+}
+
+// UpdateByID updates the document with given id with update, and remove the cache.
+func (mm *Model) UpdateByID(ctx context.Context, key string, id interface{}, update interface{},
+	opts ...*mopt.UpdateOptions) (*mongo.UpdateResult, error) {
+	res, err := mm.Model.UpdateByID(ctx, id, update, opts...)
+	if err != nil {
+		return nil, err
+	}
+
+	if err = mm.DelCache(ctx, key); err != nil {
+		return nil, err
+	}
+
+	return res, nil
+}
+
+// UpdateByIDNoCache updates the document with given id with update.
+func (mm *Model) UpdateByIDNoCache(ctx context.Context, id interface{}, update interface{},
+	opts ...*mopt.UpdateOptions) (*mongo.UpdateResult, error) {
+	return mm.Model.UpdateByID(ctx, id, update, opts...)
+}
+
+// UpdateMany updates the documents that match filter with update, and remove the cache.
+func (mm *Model) UpdateMany(ctx context.Context, keys []string, filter interface{}, update interface{},
+	opts ...*mopt.UpdateOptions) (*mongo.UpdateResult, error) {
+	res, err := mm.Model.UpdateMany(ctx, filter, update, opts...)
+	if err != nil {
+		return nil, err
+	}
+
+	if err = mm.DelCache(ctx, keys...); err != nil {
+		return nil, err
+	}
+
+	return res, nil
+}
+
+// UpdateManyNoCache updates the documents that match filter with update.
+func (mm *Model) UpdateManyNoCache(ctx context.Context, filter interface{}, update interface{},
+	opts ...*mopt.UpdateOptions) (*mongo.UpdateResult, error) {
+	return mm.Model.UpdateMany(ctx, filter, update, opts...)
+}
+
+// UpdateOne updates the first document that matches filter with update, and remove the cache.
+func (mm *Model) UpdateOne(ctx context.Context, key string, filter interface{}, update interface{},
+	opts ...*mopt.UpdateOptions) (*mongo.UpdateResult, error) {
+	res, err := mm.Model.UpdateOne(ctx, filter, update, opts...)
+	if err != nil {
+		return nil, err
+	}
+
+	if err = mm.DelCache(ctx, key); err != nil {
+		return nil, err
+	}
+
+	return res, nil
+}
+
+// UpdateOneNoCache updates the first document that matches filter with update.
+func (mm *Model) UpdateOneNoCache(ctx context.Context, filter interface{}, update interface{},
+	opts ...*mopt.UpdateOptions) (*mongo.UpdateResult, error) {
+	return mm.Model.UpdateOne(ctx, filter, update, opts...)
+}

core/stores/monc/cachedmodel_test.go

@@ -0,0 +1,581 @@
+package monc
+
+import (
+	"context"
+	"errors"
+	"sync/atomic"
+	"testing"
+
+	"github.com/alicebob/miniredis/v2"
+	"github.com/stretchr/testify/assert"
+	"github.com/zeromicro/go-zero/core/stores/cache"
+	"github.com/zeromicro/go-zero/core/stores/mon"
+	"github.com/zeromicro/go-zero/core/stores/redis"
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/mongo/integration/mtest"
+)
+
+func TestNewModel(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		_, err := newModel("foo", mt.DB.Name(), mt.Coll.Name(), nil)
+		assert.NotNil(mt, err)
+	})
+}
+
+func TestModel_DelCache(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		m := createModel(t, mt)
+		assert.Nil(t, m.cache.Set("foo", "bar"))
+		assert.Nil(t, m.cache.Set("bar", "baz"))
+		assert.Nil(t, m.DelCache(context.Background(), "foo", "bar"))
+		var v string
+		assert.True(t, m.cache.IsNotFound(m.cache.Get("foo", &v)))
+		assert.True(t, m.cache.IsNotFound(m.cache.Get("bar", &v)))
+	})
+}
+
+func TestModel_DeleteOne(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
+		m := createModel(t, mt)
+		assert.Nil(t, m.cache.Set("foo", "bar"))
+		val, err := m.DeleteOne(context.Background(), "foo", bson.D{{Key: "foo", Value: "bar"}})
+		assert.Nil(t, err)
+		assert.Equal(t, int64(1), val)
+		var v string
+		assert.True(t, m.cache.IsNotFound(m.cache.Get("foo", &v)))
+		_, err = m.DeleteOne(context.Background(), "foo", bson.D{{Key: "foo", Value: "bar"}})
+		assert.NotNil(t, err)
+
+		m.cache = mockedCache{m.cache}
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
+		_, err = m.DeleteOne(context.Background(), "foo", bson.D{{Key: "foo", Value: "bar"}})
+		assert.Equal(t, errMocked, err)
+	})
+}
+
+func TestModel_DeleteOneNoCache(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
+		m := createModel(t, mt)
+		assert.Nil(t, m.cache.Set("foo", "bar"))
+		val, err := m.DeleteOneNoCache(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
+		assert.Nil(t, err)
+		assert.Equal(t, int64(1), val)
+		var v string
+		assert.Nil(t, m.cache.Get("foo", &v))
+	})
+}
+
+func TestModel_FindOne(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		resp := mtest.CreateCursorResponse(
+			1,
+			"DBName.CollectionName",
+			mtest.FirstBatch,
+			bson.D{
+				{Key: "foo", Value: "bar"},
+			})
+		mt.AddMockResponses(resp)
+		m := createModel(t, mt)
+		var v struct {
+			Foo string `bson:"foo"`
+		}
+		assert.Nil(t, m.FindOne(context.Background(), "foo", &v, bson.D{}))
+		assert.Equal(t, "bar", v.Foo)
+		assert.Nil(t, m.cache.Set("foo", "bar"))
+	})
+}
+
+func TestModel_FindOneNoCache(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		resp := mtest.CreateCursorResponse(
+			1,
+			"DBName.CollectionName",
+			mtest.FirstBatch,
+			bson.D{
+				{Key: "foo", Value: "bar"},
+			})
+		mt.AddMockResponses(resp)
+		m := createModel(t, mt)
+		v := struct {
+			Foo string `bson:"foo"`
+		}{}
+		assert.Nil(t, m.FindOneNoCache(context.Background(), &v, bson.D{}))
+		assert.Equal(t, "bar", v.Foo)
+	})
+}
+
+func TestModel_FindOneAndDelete(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
+			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
+		}...))
+		m := createModel(t, mt)
+		assert.Nil(t, m.cache.Set("foo", "bar"))
+		v := struct {
+			Foo string `bson:"foo"`
+		}{}
+		assert.Nil(t, m.FindOneAndDelete(context.Background(), "foo", &v, bson.D{}))
+		assert.Equal(t, "bar", v.Foo)
+		assert.True(t, m.cache.IsNotFound(m.cache.Get("foo", &v)))
+		assert.NotNil(t, m.FindOneAndDelete(context.Background(), "foo", &v, bson.D{}))
+
+		m.cache = mockedCache{m.cache}
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
+			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
+		}...))
+		assert.Equal(t, errMocked, m.FindOneAndDelete(context.Background(), "foo", &v, bson.D{}))
+	})
+}
+
+func TestModel_FindOneAndDeleteNoCache(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
+			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
+		}...))
+		m := createModel(t, mt)
+		v := struct {
+			Foo string `bson:"foo"`
+		}{}
+		assert.Nil(t, m.FindOneAndDeleteNoCache(context.Background(), &v, bson.D{}))
+		assert.Equal(t, "bar", v.Foo)
+	})
+}
+
+func TestModel_FindOneAndReplace(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
+			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
+		}...))
+		m := createModel(t, mt)
+		assert.Nil(t, m.cache.Set("foo", "bar"))
+		v := struct {
+			Foo string `bson:"foo"`
+		}{}
+		assert.Nil(t, m.FindOneAndReplace(context.Background(), "foo", &v, bson.D{}, bson.D{
+			{Key: "name", Value: "Mary"},
+		}))
+		assert.Equal(t, "bar", v.Foo)
+		assert.True(t, m.cache.IsNotFound(m.cache.Get("foo", &v)))
+		assert.NotNil(t, m.FindOneAndReplace(context.Background(), "foo", &v, bson.D{}, bson.D{
+			{Key: "name", Value: "Mary"},
+		}))
+
+		m.cache = mockedCache{m.cache}
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
+			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
+		}...))
+		assert.Equal(t, errMocked, m.FindOneAndReplace(context.Background(), "foo", &v, bson.D{}, bson.D{
+			{Key: "name", Value: "Mary"},
+		}))
+	})
+}
+
+func TestModel_FindOneAndReplaceNoCache(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
+			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
+		}...))
+		m := createModel(t, mt)
+		v := struct {
+			Foo string `bson:"foo"`
+		}{}
+		assert.Nil(t, m.FindOneAndReplaceNoCache(context.Background(), &v, bson.D{}, bson.D{
+			{Key: "name", Value: "Mary"},
+		}))
+		assert.Equal(t, "bar", v.Foo)
+	})
+}
+
+func TestModel_FindOneAndUpdate(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
+			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
+		}...))
+		m := createModel(t, mt)
+		assert.Nil(t, m.cache.Set("foo", "bar"))
+		v := struct {
+			Foo string `bson:"foo"`
+		}{}
+		assert.Nil(t, m.FindOneAndUpdate(context.Background(), "foo", &v, bson.D{}, bson.D{
+			{Key: "$set", Value: bson.D{{Key: "name", Value: "Mary"}}},
+		}))
+		assert.Equal(t, "bar", v.Foo)
+		assert.True(t, m.cache.IsNotFound(m.cache.Get("foo", &v)))
+		assert.NotNil(t, m.FindOneAndUpdate(context.Background(), "foo", &v, bson.D{}, bson.D{
+			{Key: "$set", Value: bson.D{{Key: "name", Value: "Mary"}}},
+		}))
+
+		m.cache = mockedCache{m.cache}
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
+			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
+		}...))
+		assert.Equal(t, errMocked, m.FindOneAndUpdate(context.Background(), "foo", &v, bson.D{}, bson.D{
+			{Key: "$set", Value: bson.D{{Key: "name", Value: "Mary"}}},
+		}))
+	})
+}
+
+func TestModel_FindOneAndUpdateNoCache(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
+			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
+		}...))
+		m := createModel(t, mt)
+		v := struct {
+			Foo string `bson:"foo"`
+		}{}
+		assert.Nil(t, m.FindOneAndUpdateNoCache(context.Background(), &v, bson.D{}, bson.D{
+			{Key: "$set", Value: bson.D{{Key: "name", Value: "Mary"}}},
+		}))
+		assert.Equal(t, "bar", v.Foo)
+	})
+}
+
+func TestModel_GetCache(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		m := createModel(t, mt)
+		assert.NotNil(t, m.cache)
+		assert.Nil(t, m.cache.Set("foo", "bar"))
+		var s string
+		assert.Nil(t, m.cache.Get("foo", &s))
+		assert.Equal(t, "bar", s)
+	})
+}
+
+func TestModel_InsertOne(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
+			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
+		}...))
+		m := createModel(t, mt)
+		assert.Nil(t, m.cache.Set("foo", "bar"))
+		resp, err := m.InsertOne(context.Background(), "foo", bson.D{
+			{Key: "name", Value: "Mary"},
+		})
+		assert.Nil(t, err)
+		assert.NotNil(t, resp)
+		var v string
+		assert.True(t, m.cache.IsNotFound(m.cache.Get("foo", &v)))
+		_, err = m.InsertOne(context.Background(), "foo", bson.D{
+			{Key: "name", Value: "Mary"},
+		})
+		assert.NotNil(t, err)
+
+		m.cache = mockedCache{m.cache}
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
+			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
+		}...))
+		_, err = m.InsertOne(context.Background(), "foo", bson.D{
+			{Key: "name", Value: "Mary"},
+		})
+		assert.Equal(t, errMocked, err)
+	})
+}
+
+func TestModel_InsertOneNoCache(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
+			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
+		}...))
+		m := createModel(t, mt)
+		resp, err := m.InsertOneNoCache(context.Background(), bson.D{
+			{Key: "name", Value: "Mary"},
+		})
+		assert.Nil(t, err)
+		assert.NotNil(t, resp)
+	})
+}
+
+func TestModel_ReplaceOne(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
+			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
+		}...))
+		m := createModel(t, mt)
+		assert.Nil(t, m.cache.Set("foo", "bar"))
+		resp, err := m.ReplaceOne(context.Background(), "foo", bson.D{}, bson.D{
+			{Key: "foo", Value: "baz"},
+		})
+		assert.Nil(t, err)
+		assert.NotNil(t, resp)
+		var v string
+		assert.True(t, m.cache.IsNotFound(m.cache.Get("foo", &v)))
+		_, err = m.ReplaceOne(context.Background(), "foo", bson.D{}, bson.D{
+			{Key: "foo", Value: "baz"},
+		})
+		assert.NotNil(t, err)
+
+		m.cache = mockedCache{m.cache}
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
+			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
+		}...))
+		_, err = m.ReplaceOne(context.Background(), "foo", bson.D{}, bson.D{
+			{Key: "foo", Value: "baz"},
+		})
+		assert.Equal(t, errMocked, err)
+	})
+}
+
+func TestModel_ReplaceOneNoCache(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
+			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
+		}...))
+		m := createModel(t, mt)
+		resp, err := m.ReplaceOneNoCache(context.Background(), bson.D{}, bson.D{
+			{Key: "foo", Value: "baz"},
+		})
+		assert.Nil(t, err)
+		assert.NotNil(t, resp)
+	})
+}
+
+func TestModel_SetCache(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		m := createModel(t, mt)
+		assert.Nil(t, m.SetCache("foo", "bar"))
+		var v string
+		assert.Nil(t, m.GetCache("foo", &v))
+		assert.Equal(t, "bar", v)
+	})
+}
+
+func TestModel_UpdateByID(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
+			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
+		}...))
+		m := createModel(t, mt)
+		assert.Nil(t, m.cache.Set("foo", "bar"))
+		resp, err := m.UpdateByID(context.Background(), "foo", bson.D{}, bson.D{
+			{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
+		})
+		assert.Nil(t, err)
+		assert.NotNil(t, resp)
+		var v string
+		assert.True(t, m.cache.IsNotFound(m.cache.Get("foo", &v)))
+		_, err = m.UpdateByID(context.Background(), "foo", bson.D{}, bson.D{
+			{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
+		})
+		assert.NotNil(t, err)
+
+		m.cache = mockedCache{m.cache}
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
+			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
+		}...))
+		_, err = m.UpdateByID(context.Background(), "foo", bson.D{}, bson.D{
+			{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
+		})
+		assert.Equal(t, errMocked, err)
+	})
+}
+
+func TestModel_UpdateByIDNoCache(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
+			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
+		}...))
+		m := createModel(t, mt)
+		resp, err := m.UpdateByIDNoCache(context.Background(), bson.D{}, bson.D{
+			{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
+		})
+		assert.Nil(t, err)
+		assert.NotNil(t, resp)
+	})
+}
+
+func TestModel_UpdateMany(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
+			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
+		}...))
+		m := createModel(t, mt)
+		assert.Nil(t, m.cache.Set("foo", "bar"))
+		assert.Nil(t, m.cache.Set("bar", "baz"))
+		resp, err := m.UpdateMany(context.Background(), []string{"foo", "bar"}, bson.D{}, bson.D{
+			{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
+		})
+		assert.Nil(t, err)
+		assert.NotNil(t, resp)
+		var v string
+		assert.True(t, m.cache.IsNotFound(m.cache.Get("foo", &v)))
+		assert.True(t, m.cache.IsNotFound(m.cache.Get("bar", &v)))
+		_, err = m.UpdateMany(context.Background(), []string{"foo", "bar"}, bson.D{}, bson.D{
+			{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
+		})
+		assert.NotNil(t, err)
+
+		m.cache = mockedCache{m.cache}
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
+			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
+		}...))
+		_, err = m.UpdateMany(context.Background(), []string{"foo", "bar"}, bson.D{}, bson.D{
+			{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
+		})
+		assert.Equal(t, errMocked, err)
+	})
+}
+
+func TestModel_UpdateManyNoCache(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
+			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
+		}...))
+		m := createModel(t, mt)
+		resp, err := m.UpdateManyNoCache(context.Background(), bson.D{}, bson.D{
+			{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
+		})
+		assert.Nil(t, err)
+		assert.NotNil(t, resp)
+	})
+}
+
+func TestModel_UpdateOne(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
+			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
+		}...))
+		m := createModel(t, mt)
+		assert.Nil(t, m.cache.Set("foo", "bar"))
+		resp, err := m.UpdateOne(context.Background(), "foo", bson.D{}, bson.D{
+			{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
+		})
+		assert.Nil(t, err)
+		assert.NotNil(t, resp)
+		var v string
+		assert.True(t, m.cache.IsNotFound(m.cache.Get("foo", &v)))
+		_, err = m.UpdateOne(context.Background(), "foo", bson.D{}, bson.D{
+			{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
+		})
+		assert.NotNil(t, err)
+
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
+			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
+		}...))
+		m.cache = mockedCache{m.cache}
+		_, err = m.UpdateOne(context.Background(), "foo", bson.D{}, bson.D{
+			{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
+		})
+		assert.Equal(t, errMocked, err)
+	})
+}
+
+func TestModel_UpdateOneNoCache(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
+			{Key: "value", Value: bson.D{{Key: "foo", Value: "bar"}}},
+		}...))
+		m := createModel(t, mt)
+		resp, err := m.UpdateOneNoCache(context.Background(), bson.D{}, bson.D{
+			{Key: "$set", Value: bson.D{{Key: "foo", Value: "baz"}}},
+		})
+		assert.Nil(t, err)
+		assert.NotNil(t, resp)
+	})
+}
+
+func createModel(t *testing.T, mt *mtest.T) *Model {
+	s, err := miniredis.Run()
+	assert.Nil(t, err)
+	mon.Inject(mt.Name(), mt.Client)
+	if atomic.AddInt32(&index, 1)%2 == 0 {
+		return MustNewNodeModel(mt.Name(), mt.DB.Name(), mt.Coll.Name(), redis.New(s.Addr()))
+	} else {
+		return MustNewModel(mt.Name(), mt.DB.Name(), mt.Coll.Name(), cache.CacheConf{
+			cache.NodeConf{
+				RedisConf: redis.RedisConf{
+					Host: s.Addr(),
+					Type: redis.NodeType,
+				},
+				Weight: 100,
+			},
+		})
+	}
+}
+
+var (
+	errMocked = errors.New("mocked error")
+	index     int32
+)
+
+type mockedCache struct {
+	cache.Cache
+}
+
+func (m mockedCache) DelCtx(_ context.Context, _ ...string) error {
+	return errMocked
+}

core/stores/redis/hook.go

@@ -9,23 +9,33 @@ import (
 	"github.com/zeromicro/go-zero/core/logx"
 	"github.com/zeromicro/go-zero/core/mapping"
 	"github.com/zeromicro/go-zero/core/timex"
+	"github.com/zeromicro/go-zero/core/trace"
+	"go.opentelemetry.io/otel"
+	tracestd "go.opentelemetry.io/otel/trace"
 )
 
+// spanName is the span name of the redis calls.
+const spanName = "redis"
+
 var (
 	startTimeKey = contextKey("startTime")
-	durationHook = hook{}
+	durationHook = hook{tracer: otel.GetTracerProvider().Tracer(trace.TraceName)}
 )
 
 type (
 	contextKey string
-	hook       struct{}
+	hook       struct {
+		tracer tracestd.Tracer
+	}
 )
 
 func (h hook) BeforeProcess(ctx context.Context, _ red.Cmder) (context.Context, error) {
-	return context.WithValue(ctx, startTimeKey, timex.Now()), nil
+	return h.startSpan(context.WithValue(ctx, startTimeKey, timex.Now())), nil
 }
 
 func (h hook) AfterProcess(ctx context.Context, cmd red.Cmder) error {
+	h.endSpan(ctx)
+
 	val := ctx.Value(startTimeKey)
 	if val == nil {
 		return nil
@@ -45,10 +55,12 @@ func (h hook) AfterProcess(ctx context.Context, cmd red.Cmder) error {
 }
 
 func (h hook) BeforeProcessPipeline(ctx context.Context, _ []red.Cmder) (context.Context, error) {
-	return context.WithValue(ctx, startTimeKey, timex.Now()), nil
+	return h.startSpan(context.WithValue(ctx, startTimeKey, timex.Now())), nil
 }
 
 func (h hook) AfterProcessPipeline(ctx context.Context, cmds []red.Cmder) error {
+	h.endSpan(ctx)
+
 	if len(cmds) == 0 {
 		return nil
 	}
@@ -81,3 +93,12 @@ func logDuration(ctx context.Context, cmd red.Cmder, duration time.Duration) {
 	}
 	logx.WithContext(ctx).WithDuration(duration).Slowf("[REDIS] slowcall on executing: %s", buf.String())
 }
+
+func (h hook) startSpan(ctx context.Context) context.Context {
+	ctx, _ = h.tracer.Start(ctx, spanName)
+	return ctx
+}
+
+func (h hook) endSpan(ctx context.Context) {
+	tracestd.SpanFromContext(ctx).End()
+}

core/stores/redis/hook_test.go

@@ -9,9 +9,18 @@ import (
 
 	red "github.com/go-redis/redis/v8"
 	"github.com/stretchr/testify/assert"
+	ztrace "github.com/zeromicro/go-zero/core/trace"
+	tracesdk "go.opentelemetry.io/otel/trace"
 )
 
 func TestHookProcessCase1(t *testing.T) {
+	ztrace.StartAgent(ztrace.Config{
+		Name:     "go-zero-test",
+		Endpoint: "http://localhost:14268/api/traces",
+		Batcher:  "jaeger",
+		Sampler:  1.0,
+	})
+
 	writer := log.Writer()
 	var buf strings.Builder
 	log.SetOutput(&buf)
@@ -24,9 +33,17 @@ func TestHookProcessCase1(t *testing.T) {
 
 	assert.Nil(t, durationHook.AfterProcess(ctx, red.NewCmd(context.Background())))
 	assert.False(t, strings.Contains(buf.String(), "slow"))
+	assert.Equal(t, "redis", tracesdk.SpanFromContext(ctx).(interface{ Name() string }).Name())
 }
 
 func TestHookProcessCase2(t *testing.T) {
+	ztrace.StartAgent(ztrace.Config{
+		Name:     "go-zero-test",
+		Endpoint: "http://localhost:14268/api/traces",
+		Batcher:  "jaeger",
+		Sampler:  1.0,
+	})
+
 	writer := log.Writer()
 	var buf strings.Builder
 	log.SetOutput(&buf)
@@ -36,11 +53,14 @@ func TestHookProcessCase2(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+	assert.Equal(t, "redis", tracesdk.SpanFromContext(ctx).(interface{ Name() string }).Name())
 
 	time.Sleep(slowThreshold.Load() + time.Millisecond)
 
 	assert.Nil(t, durationHook.AfterProcess(ctx, red.NewCmd(context.Background(), "foo", "bar")))
 	assert.True(t, strings.Contains(buf.String(), "slow"))
+	assert.True(t, strings.Contains(buf.String(), "trace"))
+	assert.True(t, strings.Contains(buf.String(), "span"))
 }
 
 func TestHookProcessCase3(t *testing.T) {
@@ -74,6 +94,7 @@ func TestHookProcessPipelineCase1(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+	assert.Equal(t, "redis", tracesdk.SpanFromContext(ctx).(interface{ Name() string }).Name())
 
 	assert.Nil(t, durationHook.AfterProcessPipeline(ctx, []red.Cmder{
 		red.NewCmd(context.Background()),
@@ -82,6 +103,13 @@ func TestHookProcessPipelineCase1(t *testing.T) {
 }
 
 func TestHookProcessPipelineCase2(t *testing.T) {
+	ztrace.StartAgent(ztrace.Config{
+		Name:     "go-zero-test",
+		Endpoint: "http://localhost:14268/api/traces",
+		Batcher:  "jaeger",
+		Sampler:  1.0,
+	})
+
 	writer := log.Writer()
 	var buf strings.Builder
 	log.SetOutput(&buf)
@@ -91,6 +119,7 @@ func TestHookProcessPipelineCase2(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+	assert.Equal(t, "redis", tracesdk.SpanFromContext(ctx).(interface{ Name() string }).Name())
 
 	time.Sleep(slowThreshold.Load() + time.Millisecond)
 
@@ -98,6 +127,8 @@ func TestHookProcessPipelineCase2(t *testing.T) {
 		red.NewCmd(context.Background(), "foo", "bar"),
 	}))
 	assert.True(t, strings.Contains(buf.String(), "slow"))
+	assert.True(t, strings.Contains(buf.String(), "trace"))
+	assert.True(t, strings.Contains(buf.String(), "span"))
 }
 
 func TestHookProcessPipelineCase3(t *testing.T) {

core/stores/redis/redis.go

@@ -640,6 +640,29 @@ func (s *Redis) GetBitCtx(ctx context.Context, key string, offset int64) (val in
 	return
 }
 
+// GetSet is the implementation of redis getset command.
+func (s *Redis) GetSet(key, value string) (string, error) {
+	return s.GetSetCtx(context.Background(), key, value)
+}
+
+// GetSetCtx is the implementation of redis getset command.
+func (s *Redis) GetSetCtx(ctx context.Context, key, value string) (val string, err error) {
+	err = s.brk.DoWithAcceptable(func() error {
+		conn, err := getRedis(s)
+		if err != nil {
+			return err
+		}
+
+		if val, err = conn.GetSet(ctx, key, value).Result(); err == red.Nil {
+			return nil
+		}
+
+		return err
+	}, acceptable)
+
+	return
+}
+
 // Hdel is the implementation of redis hdel command.
 func (s *Redis) Hdel(key string, fields ...string) (bool, error) {
 	return s.HdelCtx(context.Background(), key, fields...)
@@ -658,7 +681,7 @@ func (s *Redis) HdelCtx(ctx context.Context, key string, fields ...string) (val
 			return err
 		}
 
-		val = v == 1
+		val = v >= 1
 		return nil
 	}, acceptable)
 
@@ -1196,7 +1219,7 @@ func (s *Redis) PfaddCtx(ctx context.Context, key string, values ...interface{})
 			return err
 		}
 
-		val = v == 1
+		val = v >= 1
 		return nil
 	}, acceptable)
 
@@ -1275,8 +1298,9 @@ func (s *Redis) Pipelined(fn func(Pipeliner) error) error {
 }
 
 // PipelinedCtx lets fn execute pipelined commands.
-func (s *Redis) PipelinedCtx(ctx context.Context, fn func(Pipeliner) error) (err error) {
-	err = s.brk.DoWithAcceptable(func() error {
+// Results need to be retrieved by calling Pipeline.Exec()
+func (s *Redis) PipelinedCtx(ctx context.Context, fn func(Pipeliner) error) error {
+	return s.brk.DoWithAcceptable(func() error {
 		conn, err := getRedis(s)
 		if err != nil {
 			return err
@@ -1285,8 +1309,6 @@ func (s *Redis) PipelinedCtx(ctx context.Context, fn func(Pipeliner) error) (err
 		_, err = conn.Pipelined(ctx, fn)
 		return err
 	}, acceptable)
-
-	return
 }
 
 // Rpop is the implementation of redis rpop command.
@@ -1381,21 +1403,28 @@ func (s *Redis) ScanCtx(ctx context.Context, cursor uint64, match string, count
 }
 
 // SetBit is the implementation of redis setbit command.
-func (s *Redis) SetBit(key string, offset int64, value int) error {
+func (s *Redis) SetBit(key string, offset int64, value int) (int, error) {
 	return s.SetBitCtx(context.Background(), key, offset, value)
 }
 
 // SetBitCtx is the implementation of redis setbit command.
-func (s *Redis) SetBitCtx(ctx context.Context, key string, offset int64, value int) error {
-	return s.brk.DoWithAcceptable(func() error {
+func (s *Redis) SetBitCtx(ctx context.Context, key string, offset int64, value int) (val int, err error) {
+	err = s.brk.DoWithAcceptable(func() error {
 		conn, err := getRedis(s)
 		if err != nil {
 			return err
 		}
 
-		_, err = conn.SetBit(ctx, key, offset, value).Result()
-		return err
+		v, err := conn.SetBit(ctx, key, offset, value).Result()
+		if err != nil {
+			return err
+		}
+
+		val = int(v)
+		return nil
 	}, acceptable)
+
+	return
 }
 
 // Sscan is the implementation of redis sscan command.

core/stores/redis/redis_test.go

@@ -387,30 +387,33 @@ func TestRedis_Mget(t *testing.T) {
 
 func TestRedis_SetBit(t *testing.T) {
 	runOnRedis(t, func(client *Redis) {
-		err := New(client.Addr, badType()).SetBit("key", 1, 1)
+		_, err := New(client.Addr, badType()).SetBit("key", 1, 1)
 		assert.NotNil(t, err)
-		err = client.SetBit("key", 1, 1)
+		val, err := client.SetBit("key", 1, 1)
 		assert.Nil(t, err)
+		assert.Equal(t, 0, val)
 	})
 }
 
 func TestRedis_GetBit(t *testing.T) {
 	runOnRedis(t, func(client *Redis) {
-		err := client.SetBit("key", 2, 1)
+		val, err := client.SetBit("key", 2, 1)
 		assert.Nil(t, err)
+		assert.Equal(t, 0, val)
 		_, err = New(client.Addr, badType()).GetBit("key", 2)
 		assert.NotNil(t, err)
-		val, err := client.GetBit("key", 2)
+		v, err := client.GetBit("key", 2)
 		assert.Nil(t, err)
-		assert.Equal(t, 1, val)
+		assert.Equal(t, 1, v)
 	})
 }
 
 func TestRedis_BitCount(t *testing.T) {
 	runOnRedis(t, func(client *Redis) {
 		for i := 0; i < 11; i++ {
-			err := client.SetBit("key", int64(i), 1)
+			val, err := client.SetBit("key", int64(i), 1)
 			assert.Nil(t, err)
+			assert.Equal(t, 0, val)
 		}
 
 		_, err := New(client.Addr, badType()).BitCount("key", 0, -1)
@@ -701,6 +704,28 @@ func TestRedis_Set(t *testing.T) {
 	})
 }
 
+func TestRedis_GetSet(t *testing.T) {
+	runOnRedis(t, func(client *Redis) {
+		_, err := New(client.Addr, badType()).GetSet("hello", "world")
+		assert.NotNil(t, err)
+		val, err := client.GetSet("hello", "world")
+		assert.Nil(t, err)
+		assert.Equal(t, "", val)
+		val, err = client.Get("hello")
+		assert.Nil(t, err)
+		assert.Equal(t, "world", val)
+		val, err = client.GetSet("hello", "newworld")
+		assert.Nil(t, err)
+		assert.Equal(t, "world", val)
+		val, err = client.Get("hello")
+		assert.Nil(t, err)
+		assert.Equal(t, "newworld", val)
+		ret, err := client.Del("hello")
+		assert.Nil(t, err)
+		assert.Equal(t, 1, ret)
+	})
+}
+
 func TestRedis_SetGetDel(t *testing.T) {
 	runOnRedis(t, func(client *Redis) {
 		err := New(client.Addr, badType()).Set("hello", "world")

core/stores/redis/redislock.go

@@ -2,6 +2,7 @@ package redis
 
 import (
 	"math/rand"
+	"strconv"
 	"sync/atomic"
 	"time"
 
@@ -11,19 +12,26 @@ import (
 )
 
 const (
+	randomLen       = 16
+	tolerance       = 500 // milliseconds
+	millisPerSecond = 1000
+	lockCommand     = `if redis.call("GET", KEYS[1]) == ARGV[1] then
+    redis.call("SET", KEYS[1], ARGV[1], "PX", ARGV[2])
+    return "OK"
+else
+    return redis.call("SET", KEYS[1], ARGV[1], "NX", "PX", ARGV[2])
+end`
 	delCommand = `if redis.call("GET", KEYS[1]) == ARGV[1] then
     return redis.call("DEL", KEYS[1])
 else
     return 0
 end`
-	randomLen = 16
 )
 
 // A RedisLock is a redis lock.
 type RedisLock struct {
 	store   *Redis
 	seconds uint32
-	count   int32
 	key     string
 	id      string
 }
@@ -43,35 +51,30 @@ func NewRedisLock(store *Redis, key string) *RedisLock {
 
 // Acquire acquires the lock.
 func (rl *RedisLock) Acquire() (bool, error) {
-	newCount := atomic.AddInt32(&rl.count, 1)
-	if newCount > 1 {
+	seconds := atomic.LoadUint32(&rl.seconds)
+	resp, err := rl.store.Eval(lockCommand, []string{rl.key}, []string{
+		rl.id, strconv.Itoa(int(seconds)*millisPerSecond + tolerance),
+	})
+	if err == red.Nil {
+		return false, nil
+	} else if err != nil {
+		logx.Errorf("Error on acquiring lock for %s, %s", rl.key, err.Error())
+		return false, err
+	} else if resp == nil {
+		return false, nil
+	}
+
+	reply, ok := resp.(string)
+	if ok && reply == "OK" {
 		return true, nil
 	}
 
-	seconds := atomic.LoadUint32(&rl.seconds)
-	ok, err := rl.store.SetnxEx(rl.key, rl.id, int(seconds+1)) // +1s for tolerance
-	if err == red.Nil {
-		atomic.AddInt32(&rl.count, -1)
-		return false, nil
-	} else if err != nil {
-		atomic.AddInt32(&rl.count, -1)
-		logx.Errorf("Error on acquiring lock for %s, %s", rl.key, err.Error())
-		return false, err
-	} else if !ok {
-		atomic.AddInt32(&rl.count, -1)
-		return false, nil
-	}
-
-	return true, nil
+	logx.Errorf("Unknown reply when acquiring lock for %s: %v", rl.key, resp)
+	return false, nil
 }
 
 // Release releases the lock.
 func (rl *RedisLock) Release() (bool, error) {
-	newCount := atomic.AddInt32(&rl.count, -1)
-	if newCount > 0 {
-		return true, nil
-	}
-
 	resp, err := rl.store.Eval(delCommand, []string{rl.key}, []string{rl.id})
 	if err != nil {
 		return false, err

core/stores/redis/redislock_test.go

@@ -29,25 +29,5 @@ func TestRedisLock(t *testing.T) {
 		endAcquire, err := secondLock.Acquire()
 		assert.Nil(t, err)
 		assert.True(t, endAcquire)
-
-		endAcquire, err = secondLock.Acquire()
-		assert.Nil(t, err)
-		assert.True(t, endAcquire)
-
-		release, err = secondLock.Release()
-		assert.Nil(t, err)
-		assert.True(t, release)
-
-		againAcquire, err = firstLock.Acquire()
-		assert.Nil(t, err)
-		assert.False(t, againAcquire)
-
-		release, err = secondLock.Release()
-		assert.Nil(t, err)
-		assert.True(t, release)
-
-		firstAcquire, err = firstLock.Acquire()
-		assert.Nil(t, err)
-		assert.True(t, firstAcquire)
 	})
 }

core/stores/sqlx/sqlconn.go

@@ -6,8 +6,14 @@ import (
 
 	"github.com/zeromicro/go-zero/core/breaker"
 	"github.com/zeromicro/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/core/trace"
+	"go.opentelemetry.io/otel"
+	tracesdk "go.opentelemetry.io/otel/trace"
 )
 
+// spanName is used to identify the span name for the SQL execution.
+const spanName = "sql"
+
 // ErrNotFound is an alias of sql.ErrNoRows
 var ErrNotFound = sql.ErrNoRows
 
@@ -134,6 +140,9 @@ func (db *commonSqlConn) Exec(q string, args ...interface{}) (result sql.Result,
 
 func (db *commonSqlConn) ExecCtx(ctx context.Context, q string, args ...interface{}) (
 	result sql.Result, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	err = db.brk.DoWithAcceptable(func() error {
 		var conn *sql.DB
 		conn, err = db.connProv()
@@ -154,6 +163,9 @@ func (db *commonSqlConn) Prepare(query string) (stmt StmtSession, err error) {
 }
 
 func (db *commonSqlConn) PrepareCtx(ctx context.Context, query string) (stmt StmtSession, err error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	err = db.brk.DoWithAcceptable(func() error {
 		var conn *sql.DB
 		conn, err = db.connProv()
@@ -183,6 +195,9 @@ func (db *commonSqlConn) QueryRow(v interface{}, q string, args ...interface{})
 
 func (db *commonSqlConn) QueryRowCtx(ctx context.Context, v interface{}, q string,
 	args ...interface{}) error {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	return db.queryRows(ctx, func(rows *sql.Rows) error {
 		return unmarshalRow(v, rows, true)
 	}, q, args...)
@@ -194,6 +209,9 @@ func (db *commonSqlConn) QueryRowPartial(v interface{}, q string, args ...interf
 
 func (db *commonSqlConn) QueryRowPartialCtx(ctx context.Context, v interface{},
 	q string, args ...interface{}) error {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	return db.queryRows(ctx, func(rows *sql.Rows) error {
 		return unmarshalRow(v, rows, false)
 	}, q, args...)
@@ -205,6 +223,9 @@ func (db *commonSqlConn) QueryRows(v interface{}, q string, args ...interface{})
 
 func (db *commonSqlConn) QueryRowsCtx(ctx context.Context, v interface{}, q string,
 	args ...interface{}) error {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	return db.queryRows(ctx, func(rows *sql.Rows) error {
 		return unmarshalRows(v, rows, true)
 	}, q, args...)
@@ -216,6 +237,9 @@ func (db *commonSqlConn) QueryRowsPartial(v interface{}, q string, args ...inter
 
 func (db *commonSqlConn) QueryRowsPartialCtx(ctx context.Context, v interface{},
 	q string, args ...interface{}) error {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	return db.queryRows(ctx, func(rows *sql.Rows) error {
 		return unmarshalRows(v, rows, false)
 	}, q, args...)
@@ -232,6 +256,9 @@ func (db *commonSqlConn) Transact(fn func(Session) error) error {
 }
 
 func (db *commonSqlConn) TransactCtx(ctx context.Context, fn func(context.Context, Session) error) error {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	return db.brk.DoWithAcceptable(func() error {
 		return transact(ctx, db, db.beginTx, fn)
 	}, db.acceptable)
@@ -248,6 +275,9 @@ func (db *commonSqlConn) acceptable(err error) bool {
 
 func (db *commonSqlConn) queryRows(ctx context.Context, scanner func(*sql.Rows) error,
 	q string, args ...interface{}) error {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	var qerr error
 	return db.brk.DoWithAcceptable(func() error {
 		conn, err := db.connProv()
@@ -274,6 +304,9 @@ func (s statement) Exec(args ...interface{}) (sql.Result, error) {
 }
 
 func (s statement) ExecCtx(ctx context.Context, args ...interface{}) (sql.Result, error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	return execStmt(ctx, s.stmt, s.query, args...)
 }
 
@@ -282,6 +315,9 @@ func (s statement) QueryRow(v interface{}, args ...interface{}) error {
 }
 
 func (s statement) QueryRowCtx(ctx context.Context, v interface{}, args ...interface{}) error {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	return queryStmt(ctx, s.stmt, func(rows *sql.Rows) error {
 		return unmarshalRow(v, rows, true)
 	}, s.query, args...)
@@ -292,6 +328,9 @@ func (s statement) QueryRowPartial(v interface{}, args ...interface{}) error {
 }
 
 func (s statement) QueryRowPartialCtx(ctx context.Context, v interface{}, args ...interface{}) error {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	return queryStmt(ctx, s.stmt, func(rows *sql.Rows) error {
 		return unmarshalRow(v, rows, false)
 	}, s.query, args...)
@@ -302,6 +341,9 @@ func (s statement) QueryRows(v interface{}, args ...interface{}) error {
 }
 
 func (s statement) QueryRowsCtx(ctx context.Context, v interface{}, args ...interface{}) error {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	return queryStmt(ctx, s.stmt, func(rows *sql.Rows) error {
 		return unmarshalRows(v, rows, true)
 	}, s.query, args...)
@@ -312,7 +354,15 @@ func (s statement) QueryRowsPartial(v interface{}, args ...interface{}) error {
 }
 
 func (s statement) QueryRowsPartialCtx(ctx context.Context, v interface{}, args ...interface{}) error {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	return queryStmt(ctx, s.stmt, func(rows *sql.Rows) error {
 		return unmarshalRows(v, rows, false)
 	}, s.query, args...)
 }
+
+func startSpan(ctx context.Context) (context.Context, tracesdk.Span) {
+	tracer := otel.GetTracerProvider().Tracer(trace.TraceName)
+	return tracer.Start(ctx, spanName)
+}

core/stores/sqlx/tx.go

@@ -31,6 +31,9 @@ func (t txSession) Exec(q string, args ...interface{}) (sql.Result, error) {
 }
 
 func (t txSession) ExecCtx(ctx context.Context, q string, args ...interface{}) (sql.Result, error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	return exec(ctx, t.Tx, q, args...)
 }
 
@@ -39,6 +42,9 @@ func (t txSession) Prepare(q string) (StmtSession, error) {
 }
 
 func (t txSession) PrepareCtx(ctx context.Context, q string) (StmtSession, error) {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	stmt, err := t.Tx.PrepareContext(ctx, q)
 	if err != nil {
 		return nil, err
@@ -55,6 +61,9 @@ func (t txSession) QueryRow(v interface{}, q string, args ...interface{}) error
 }
 
 func (t txSession) QueryRowCtx(ctx context.Context, v interface{}, q string, args ...interface{}) error {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	return query(ctx, t.Tx, func(rows *sql.Rows) error {
 		return unmarshalRow(v, rows, true)
 	}, q, args...)
@@ -66,6 +75,9 @@ func (t txSession) QueryRowPartial(v interface{}, q string, args ...interface{})
 
 func (t txSession) QueryRowPartialCtx(ctx context.Context, v interface{}, q string,
 	args ...interface{}) error {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	return query(ctx, t.Tx, func(rows *sql.Rows) error {
 		return unmarshalRow(v, rows, false)
 	}, q, args...)
@@ -76,6 +88,9 @@ func (t txSession) QueryRows(v interface{}, q string, args ...interface{}) error
 }
 
 func (t txSession) QueryRowsCtx(ctx context.Context, v interface{}, q string, args ...interface{}) error {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	return query(ctx, t.Tx, func(rows *sql.Rows) error {
 		return unmarshalRows(v, rows, true)
 	}, q, args...)
@@ -87,6 +102,9 @@ func (t txSession) QueryRowsPartial(v interface{}, q string, args ...interface{}
 
 func (t txSession) QueryRowsPartialCtx(ctx context.Context, v interface{}, q string,
 	args ...interface{}) error {
+	ctx, span := startSpan(ctx)
+	defer span.End()
+
 	return query(ctx, t.Tx, func(rows *sql.Rows) error {
 		return unmarshalRows(v, rows, false)
 	}, q, args...)

core/stores/sqlx/utils.go

@@ -2,6 +2,7 @@ package sqlx
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"strconv"
 	"strings"
@@ -10,6 +11,8 @@ import (
 	"github.com/zeromicro/go-zero/core/mapping"
 )
 
+var errUnbalancedEscape = errors.New("no char after escape char")
+
 func desensitize(datasource string) string {
 	// remove account
 	pos := strings.LastIndex(datasource, "@")
@@ -75,6 +78,7 @@ func format(query string, args ...interface{}) (string, error) {
 					break
 				}
 			}
+
 			if j > i+1 {
 				index, err := strconv.Atoi(query[i+1 : j])
 				if err != nil {
@@ -85,7 +89,7 @@ func format(query string, args ...interface{}) (string, error) {
 				if index > argIndex {
 					argIndex = index
 				}
-				
+
 				index--
 				if index < 0 || numArgs <= index {
 					return "", fmt.Errorf("error: wrong index %d in sql", index)
@@ -94,6 +98,25 @@ func format(query string, args ...interface{}) (string, error) {
 				writeValue(&b, args[index])
 				i = j - 1
 			}
+		case '\'', '"', '`':
+			b.WriteByte(ch)
+
+			for j := i + 1; j < bytes; j++ {
+				cur := query[j]
+				b.WriteByte(cur)
+
+				if cur == '\\' {
+					j++
+					if j >= bytes {
+						return "", errUnbalancedEscape
+					}
+
+					b.WriteByte(query[j])
+				} else if cur == ch {
+					i = j
+					break
+				}
+			}
 		default:
 			b.WriteByte(ch)
 		}

core/stores/sqlx/utils_test.go

@@ -97,6 +97,30 @@ func TestFormat(t *testing.T) {
 			args:   []interface{}{"133", false},
 			hasErr: true,
 		},
+		{
+			name:   "select with date",
+			query:  "select * from user where date='2006-01-02 15:04:05' and name=:1",
+			args:   []interface{}{"foo"},
+			expect: "select * from user where date='2006-01-02 15:04:05' and name='foo'",
+		},
+		{
+			name:   "select with date and escape",
+			query:  `select * from user where date=' 2006-01-02 15:04:05 \'' and name=:1`,
+			args:   []interface{}{"foo"},
+			expect: `select * from user where date=' 2006-01-02 15:04:05 \'' and name='foo'`,
+		},
+		{
+			name:   "select with date and bad arg",
+			query:  `select * from user where date='2006-01-02 15:04:05 \'' and name=:a`,
+			args:   []interface{}{"foo"},
+			hasErr: true,
+		},
+		{
+			name:   "select with date and escape error",
+			query:  `select * from user where date='2006-01-02 15:04:05 \`,
+			args:   []interface{}{"foo"},
+			hasErr: true,
+		},
 	}
 
 	for _, test := range tests {
@@ -108,6 +132,7 @@ func TestFormat(t *testing.T) {
 			if test.hasErr {
 				assert.NotNil(t, err)
 			} else {
+				assert.Nil(t, err)
 				assert.Equal(t, test.expect, actual)
 			}
 		})

core/syncx/resourcemanager.go

@@ -68,3 +68,10 @@ func (manager *ResourceManager) GetResource(key string, create func() (io.Closer
 
 	return val.(io.Closer), nil
 }
+
+// Inject injects the resource associated with given key.
+func (manager *ResourceManager) Inject(key string, resource io.Closer) {
+	manager.lock.Lock()
+	manager.resources[key] = resource
+	manager.lock.Unlock()
+}

core/syncx/resourcemanager_test.go

@@ -47,6 +47,8 @@ func TestResourceManager_GetResourceError(t *testing.T) {
 
 func TestResourceManager_Close(t *testing.T) {
 	manager := NewResourceManager()
+	defer manager.Close()
+
 	for i := 0; i < 10; i++ {
 		_, err := manager.GetResource("key", func() (io.Closer, error) {
 			return nil, errors.New("fail")
@@ -61,6 +63,8 @@ func TestResourceManager_Close(t *testing.T) {
 
 func TestResourceManager_UseAfterClose(t *testing.T) {
 	manager := NewResourceManager()
+	defer manager.Close()
+
 	_, err := manager.GetResource("key", func() (io.Closer, error) {
 		return nil, errors.New("fail")
 	})
@@ -72,3 +76,18 @@ func TestResourceManager_UseAfterClose(t *testing.T) {
 		assert.NotNil(t, err)
 	}
 }
+
+func TestResourceManager_Inject(t *testing.T) {
+	manager := NewResourceManager()
+	defer manager.Close()
+
+	manager.Inject("key", &dummyResource{
+		age: 10,
+	})
+
+	val, err := manager.GetResource("key", func() (io.Closer, error) {
+		return nil, nil
+	})
+	assert.Nil(t, err)
+	assert.Equal(t, 10, val.(*dummyResource).age)
+}

go.mod

@@ -7,6 +7,7 @@ require (
 	github.com/DATA-DOG/go-sqlmock v1.5.0
 	github.com/alicebob/miniredis/v2 v2.17.0
 	github.com/globalsign/mgo v0.0.0-20181015135952-eeefdecb41b8
+	github.com/go-redis/redis/v8 v8.11.4
 	github.com/go-sql-driver/mysql v1.6.0
 	github.com/golang-jwt/jwt/v4 v4.2.0
 	github.com/golang/mock v1.6.0
@@ -19,6 +20,7 @@ require (
 	github.com/stretchr/testify v1.7.0
 	go.etcd.io/etcd/api/v3 v3.5.2
 	go.etcd.io/etcd/client/v3 v3.5.2
+	go.mongodb.org/mongo-driver v1.9.0
 	go.opentelemetry.io/otel v1.3.0
 	go.opentelemetry.io/otel/exporters/jaeger v1.3.0
 	go.opentelemetry.io/otel/exporters/zipkin v1.3.0
@@ -26,9 +28,10 @@ require (
 	go.opentelemetry.io/otel/trace v1.3.0
 	go.uber.org/automaxprocs v1.4.0
 	go.uber.org/goleak v1.1.12
+	golang.org/x/sys v0.0.0-20220422013727-9388b58f7150
 	golang.org/x/time v0.0.0-20211116232009-f0f3c7e86c11
-	google.golang.org/grpc v1.44.0
-	google.golang.org/protobuf v1.27.1
+	google.golang.org/grpc v1.46.0
+	google.golang.org/protobuf v1.28.0
 	gopkg.in/cheggaaa/pb.v1 v1.0.28
 	gopkg.in/h2non/gock.v1 v1.1.2
 	gopkg.in/yaml.v2 v2.4.0
@@ -41,7 +44,6 @@ require (
 require (
 	github.com/fatih/color v1.10.0 // indirect
 	github.com/go-logr/stdr v1.2.2 // indirect
-	github.com/go-redis/redis/v8 v8.11.4
 	github.com/google/gofuzz v1.2.0 // indirect
 	github.com/json-iterator/go v1.1.12 // indirect
 	github.com/mattn/go-runewidth v0.0.13 // indirect
@@ -52,10 +54,9 @@ require (
 	go.uber.org/atomic v1.9.0 // indirect
 	go.uber.org/multierr v1.8.0 // indirect
 	go.uber.org/zap v1.21.0 // indirect
-	golang.org/x/net v0.0.0-20220225172249-27dd8689420f // indirect
+	golang.org/x/net v0.0.0-20220421235706-1d1ef9303861 // indirect
 	golang.org/x/oauth2 v0.0.0-20210819190943-2bc19b11175f // indirect
-	golang.org/x/sys v0.0.0-20220227234510-4e6760a101f9
 	google.golang.org/appengine v1.6.7 // indirect
-	google.golang.org/genproto v0.0.0-20220228195345-15d65a4533f7 // indirect
+	google.golang.org/genproto v0.0.0-20220422154200-b37d22cd5731 // indirect
 	k8s.io/klog/v2 v2.40.1 // indirect
 )

go.sum

@@ -85,6 +85,7 @@ github.com/cncf/udpa/go v0.0.0-20201120205902-5459f2c99403/go.mod h1:WmhPx2Nbnht
 github.com/cncf/udpa/go v0.0.0-20210930031921-04548b0d99d4/go.mod h1:6pvJx4me5XPnfI9Z40ddWsdw2W/uZgQLFXToKeRcDiI=
 github.com/cncf/xds/go v0.0.0-20210805033703-aa0b78936158/go.mod h1:eXthEFrGJvWHgFFCl3hGmgk+/aYT6PnTQLykKQRLhEs=
 github.com/cncf/xds/go v0.0.0-20210922020428-25de7278fc84/go.mod h1:eXthEFrGJvWHgFFCl3hGmgk+/aYT6PnTQLykKQRLhEs=
+github.com/cncf/xds/go v0.0.0-20211001041855-01bcc9b48dfe/go.mod h1:eXthEFrGJvWHgFFCl3hGmgk+/aYT6PnTQLykKQRLhEs=
 github.com/cncf/xds/go v0.0.0-20211011173535-cb28da3451f1/go.mod h1:eXthEFrGJvWHgFFCl3hGmgk+/aYT6PnTQLykKQRLhEs=
 github.com/coreos/go-semver v0.3.0 h1:wkHLiw0WNATZnSG7epLsujiMCgPAc9xhjJ4tgnAxmfM=
 github.com/coreos/go-semver v0.3.0/go.mod h1:nnelYz7RCh+5ahJtPPxZlU+153eP4D4r3EedlOD2RNk=
@@ -112,6 +113,7 @@ github.com/envoyproxy/go-control-plane v0.9.4/go.mod h1:6rpuAdCZL397s3pYoYcLgu1m
 github.com/envoyproxy/go-control-plane v0.9.9-0.20201210154907-fd9021fe5dad/go.mod h1:cXg6YxExXjJnVBQHBLXeUAgxn2UodCpnH306RInaBQk=
 github.com/envoyproxy/go-control-plane v0.9.9-0.20210217033140-668b12f5399d/go.mod h1:cXg6YxExXjJnVBQHBLXeUAgxn2UodCpnH306RInaBQk=
 github.com/envoyproxy/go-control-plane v0.9.10-0.20210907150352-cf90f659a021/go.mod h1:AFq3mo9L8Lqqiid3OhADV3RfLJnjiw63cSpi+fDTRC0=
+github.com/envoyproxy/go-control-plane v0.10.2-0.20220325020618-49ff273808a1/go.mod h1:KJwIaB5Mv44NWtYuAOFCVOjcI94vtpEz2JU/D2v6IjE=
 github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c=
 github.com/evanphx/json-patch v4.9.0+incompatible/go.mod h1:50XU6AFN0ol/bzJsmQLiYLvXMP4fmwYFNcr97nuDLSk=
 github.com/fatih/color v1.10.0 h1:s36xzo75JdqLaaWoiEHk767eHiwo0598uUxyfiPkDsg=
@@ -157,6 +159,7 @@ github.com/go-redis/redis/v8 v8.11.4/go.mod h1:2Z2wHZXdQpCDXEGzqMockDpNyYvi2l4Px
 github.com/go-sql-driver/mysql v1.4.0/go.mod h1:zAC/RDZ24gD3HViQzih4MyKcchzm+sOG5ZlKdlhCg5w=
 github.com/go-sql-driver/mysql v1.6.0 h1:BCTh4TKNUYmOmMUcQ3IipzF5prigylS7XXjEkfCHuOE=
 github.com/go-sql-driver/mysql v1.6.0/go.mod h1:DCzpHaOWr8IXmIStZouvnhqoel9Qv2LBy8hT2VhHyBg=
+github.com/go-stack/stack v1.8.0 h1:5SgMzNM5HxrEjV0ww2lTmX6E2Izsfxas4+YHWRs3Lsk=
 github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY=
 github.com/go-task/slim-sprig v0.0.0-20210107165309-348f09dbbbc0/go.mod h1:fyg7847qk6SyHyPtNmDHnmrv/HOrqktSC+C9fM+CJOE=
 github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
@@ -196,6 +199,8 @@ github.com/golang/protobuf v1.4.3/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw
 github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
 github.com/golang/protobuf v1.5.2 h1:ROPKBNFfQgOUMifHyP+KYbvpjbdoFNs+aK7DXlji0Tw=
 github.com/golang/protobuf v1.5.2/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY=
+github.com/golang/snappy v0.0.1/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
+github.com/golang/snappy v0.0.4 h1:yAGX7huGHXlcLOEtBnF4w7FQwA26wojNCwOYAEhLjQM=
 github.com/golang/snappy v0.0.4/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
 github.com/google/btree v0.0.0-20180813153112-4030bb1f1f0c/go.mod h1:lNA+9X1NB3Zf8V7Ke586lFgjr2dZNuvo3lPJSGZ5JPQ=
 github.com/google/btree v1.0.0/go.mod h1:lNA+9X1NB3Zf8V7Ke586lFgjr2dZNuvo3lPJSGZ5JPQ=
@@ -269,6 +274,7 @@ github.com/justinas/alice v1.2.0 h1:+MHSA/vccVCF4Uq37S42jwlkvI2Xzl7zTPCN5BnZNVo=
 github.com/justinas/alice v1.2.0/go.mod h1:fN5HRH/reO/zrUflLfTN43t3vXvKzvZIENsNEe7i7qA=
 github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8=
 github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
+github.com/klauspost/compress v1.13.6 h1:P76CopJELS0TiO2mebmnzgWaajssP/EszplttgQxcgc=
 github.com/klauspost/compress v1.13.6/go.mod h1:/3/Vjq9QcHkK5uEr5lBEmyoZ1iFhe47etQ6QUkpK6sk=
 github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=
 github.com/konsorten/go-windows-terminal-sequences v1.0.3/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=
@@ -306,6 +312,7 @@ github.com/modern-go/reflect2 v0.0.0-20180701023420-4b7aa43c6742/go.mod h1:bx2lN
 github.com/modern-go/reflect2 v1.0.1/go.mod h1:bx2lNnkwVCuqBIxFjflWJWanXIb3RllmbCylyMrvgv0=
 github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M=
 github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk=
+github.com/montanaflynn/stats v0.0.0-20171201202039-1bf9dbcd8cbe/go.mod h1:wL8QJuTMNUDYhXwkmfOly8iTdp5TEcJFWZD2D7SIkUc=
 github.com/munnerz/goautoneg v0.0.0-20120707110453-a547fc61f48d/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ=
 github.com/mwitkow/go-conntrack v0.0.0-20161129095857-cc309e4a2223/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U=
 github.com/mwitkow/go-conntrack v0.0.0-20190716064945-2f068394615f/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U=
@@ -393,10 +400,17 @@ github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5
 github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY=
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+github.com/tidwall/pretty v1.0.0 h1:HsD+QiTn7sK6flMKIvNmpqz1qrpP3Ps6jOKIKMooyg4=
+github.com/tidwall/pretty v1.0.0/go.mod h1:XNkn88O1ChpSDQmQeStsy+sBenx6DDtFZJxhVysOjyk=
 github.com/urfave/cli/v2 v2.3.0/go.mod h1:LJmUH05zAU44vOAcrfzZQKsZbVcdbOG8rtL3/XcUArI=
+github.com/xdg-go/pbkdf2 v1.0.0 h1:Su7DPu48wXMwC3bs7MCNG+z4FhcyEuz5dlvchbq0B0c=
 github.com/xdg-go/pbkdf2 v1.0.0/go.mod h1:jrpuAogTd400dnrH08LKmI/xc1MbPOebTwRqcT5RDeI=
+github.com/xdg-go/scram v1.0.2 h1:akYIkZ28e6A96dkWNJQu3nmCzH3YfwMPQExUYDaRv7w=
 github.com/xdg-go/scram v1.0.2/go.mod h1:1WAq6h33pAW+iRreB34OORO2Nf7qel3VV3fjBj+hCSs=
+github.com/xdg-go/stringprep v1.0.2 h1:6iq84/ryjjeRmMJwxutI51F2GIPlP5BfTvXHeYjyhBc=
 github.com/xdg-go/stringprep v1.0.2/go.mod h1:8F9zXuvzgwmyT5DUm4GUfZGDdT3W+LCvS6+da4O5kxM=
+github.com/youmark/pkcs8 v0.0.0-20181117223130-1be2e3e5546d h1:splanxYIlg+5LfHAM6xpdFEAYOk8iySO56hMFq6uLyA=
+github.com/youmark/pkcs8 v0.0.0-20181117223130-1be2e3e5546d/go.mod h1:rHwXgn7JulP+udvsHwJoVG1YGAP6VLg4y9I5dyZdqmA=
 github.com/yuin/goldmark v1.1.25/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.1.32/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
@@ -410,6 +424,8 @@ go.etcd.io/etcd/client/pkg/v3 v3.5.2 h1:4hzqQ6hIb3blLyQ8usCU4h3NghkqcsohEQ3o3Vet
 go.etcd.io/etcd/client/pkg/v3 v3.5.2/go.mod h1:IJHfcCEKxYu1Os13ZdwCwIUTUVGYTSAM3YSwc9/Ac1g=
 go.etcd.io/etcd/client/v3 v3.5.2 h1:WdnejrUtQC4nCxK0/dLTMqKOB+U5TP/2Ya0BJL+1otA=
 go.etcd.io/etcd/client/v3 v3.5.2/go.mod h1:kOOaWFFgHygyT0WlSmL8TJiXmMysO/nNUlEsSsN6W4o=
+go.mongodb.org/mongo-driver v1.9.0 h1:f3aLGJvQmBl8d9S40IL+jEyBC6hfLPbJjv9t5hEM9ck=
+go.mongodb.org/mongo-driver v1.9.0/go.mod h1:0sQWfOeY63QTntERDJJ/0SuKK0T1uVSgKCuAROlKEPY=
 go.opencensus.io v0.21.0/go.mod h1:mSImk1erAIZhrmZN+AvHh14ztQfjbGwt4TtuofqLduU=
 go.opencensus.io v0.22.0/go.mod h1:+kGneAE2xo2IficOXnaByMWTGM9T73dGwxeWcUqIpI8=
 go.opencensus.io v0.22.2/go.mod h1:yxeiOL68Rb0Xd1ddK5vPZ/oVn4vY4Ynel7k9FzqtOIw=
@@ -449,6 +465,7 @@ golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8U
 golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
 golang.org/x/crypto v0.0.0-20201002170205-7f63de1d35b0/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
 golang.org/x/crypto v0.0.0-20201112155050-0c6587e931a9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
+golang.org/x/crypto v0.0.0-20201216223049-8b5274cf687f/go.mod h1:jdWPYTVW3xRLrWPugEBEK3UY2ZEsg3UU495nc5E+M+I=
 golang.org/x/crypto v0.0.0-20210920023735-84f357641f63 h1:kETrAMYZq6WVGPa8IIixL0CaEcIUNi+1WX7grUoi3y8=
 golang.org/x/crypto v0.0.0-20210920023735-84f357641f63/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
 golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
@@ -523,8 +540,8 @@ golang.org/x/net v0.0.0-20210405180319-a5a99cb37ef4/go.mod h1:p54w0d4576C0XHj96b
 golang.org/x/net v0.0.0-20210428140749-89ef3d95e781/go.mod h1:OJAsFXCWl8Ukc7SiCT/9KSuxbyM7479/AVlXFRxuMCk=
 golang.org/x/net v0.0.0-20210525063256-abc453219eb5/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=
 golang.org/x/net v0.0.0-20210917221730-978cfadd31cf/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=
-golang.org/x/net v0.0.0-20220225172249-27dd8689420f h1:oA4XRj0qtSt8Yo1Zms0CUlsT3KG69V2UGQWPBxujDmc=
-golang.org/x/net v0.0.0-20220225172249-27dd8689420f/go.mod h1:CfG3xpIq0wQ8r1q4Su4UZFWDARRcnwPjda9FqA0JpMk=
+golang.org/x/net v0.0.0-20220421235706-1d1ef9303861 h1:yssD99+7tqHWO5Gwh81phT+67hg+KttniBr6UnEXOY8=
+golang.org/x/net v0.0.0-20220421235706-1d1ef9303861/go.mod h1:CfG3xpIq0wQ8r1q4Su4UZFWDARRcnwPjda9FqA0JpMk=
 golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U=
 golang.org/x/oauth2 v0.0.0-20190226205417-e64efc72b421/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw=
 golang.org/x/oauth2 v0.0.0-20190604053449-0f29369cfe45/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw=
@@ -543,6 +560,7 @@ golang.org/x/sync v0.0.0-20200317015054-43a5402ce75a/go.mod h1:RxMgew5VJxzue5/jJ
 golang.org/x/sync v0.0.0-20200625203802-6e8e738ad208/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20201207232520-09787c993a3a/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.0.0-20210220032951-036812b2e83c h1:5KslGYwFpkhGh+Q16bwMP3cOontH8FOep7tGV86Y7SQ=
 golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20180905080454-ebe1bf3edb33/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
@@ -587,6 +605,7 @@ golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/sys v0.0.0-20201112073958-5cba982894dd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210112080510-489259a85091/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20210119212857-b64e53b001e4/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210124154548-22da62e12c0c/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210330210617-4fbd30eecc44/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210403161142-5e06dd20ab57/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
@@ -596,8 +615,9 @@ golang.org/x/sys v0.0.0-20210510120138-977fb7262007/go.mod h1:oPkhp1MJrh7nUepCBc
 golang.org/x/sys v0.0.0-20210603081109-ebe580a85c40/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20211216021012-1d35b9e2eb4e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.0.0-20220227234510-4e6760a101f9 h1:nhht2DYV/Sn3qOayu8lM+cU1ii9sTLUeBQwQQfUHtrs=
-golang.org/x/sys v0.0.0-20220227234510-4e6760a101f9/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20220422013727-9388b58f7150 h1:xHms4gcpe1YE7A3yIllJXP16CMAGuqwO2lX1mTyyRRc=
+golang.org/x/sys v0.0.0-20220422013727-9388b58f7150/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/term v0.0.0-20201117132131-f5c789dd3221/go.mod h1:Nr5EML6q2oocZ2LXRh80K7BxOlk5/8JxuGnuhpl+muw=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/term v0.0.0-20210927222741-03fcf44c2211 h1:JGgROgKl9N8DuW20oFS5gxc+lE67/N3FcwmBPMe7ArY=
 golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
@@ -625,6 +645,7 @@ golang.org/x/tools v0.0.0-20190312170243-e65039ee4138/go.mod h1:LCzVGOaR6xXOjkQ3
 golang.org/x/tools v0.0.0-20190425150028-36563e24a262/go.mod h1:RgjU9mgBXZiqYHBnxXauZ1Gv1EHHAz9KjViQ78xBX0Q=
 golang.org/x/tools v0.0.0-20190506145303-2d16b83fe98c/go.mod h1:RgjU9mgBXZiqYHBnxXauZ1Gv1EHHAz9KjViQ78xBX0Q=
 golang.org/x/tools v0.0.0-20190524140312-2c0ae7006135/go.mod h1:RgjU9mgBXZiqYHBnxXauZ1Gv1EHHAz9KjViQ78xBX0Q=
+golang.org/x/tools v0.0.0-20190531172133-b3315ee88b7d/go.mod h1:/rFqwRUd4F7ZHNgwSSTFct+R/Kf4OFW1sUzUTQQTgfc=
 golang.org/x/tools v0.0.0-20190606124116-d0a3d012864b/go.mod h1:/rFqwRUd4F7ZHNgwSSTFct+R/Kf4OFW1sUzUTQQTgfc=
 golang.org/x/tools v0.0.0-20190614205625-5aca471b1d59/go.mod h1:/rFqwRUd4F7ZHNgwSSTFct+R/Kf4OFW1sUzUTQQTgfc=
 golang.org/x/tools v0.0.0-20190621195816-6e04913cbbac/go.mod h1:/rFqwRUd4F7ZHNgwSSTFct+R/Kf4OFW1sUzUTQQTgfc=
@@ -724,8 +745,8 @@ google.golang.org/genproto v0.0.0-20200729003335-053ba62fc06f/go.mod h1:FWY/as6D
 google.golang.org/genproto v0.0.0-20200804131852-c06518451d9c/go.mod h1:FWY/as6DDZQgahTzZj3fqbO1CbirC29ZNUFHwi0/+no=
 google.golang.org/genproto v0.0.0-20200825200019-8632dd797987/go.mod h1:FWY/as6DDZQgahTzZj3fqbO1CbirC29ZNUFHwi0/+no=
 google.golang.org/genproto v0.0.0-20210602131652-f16073e35f0c/go.mod h1:UODoCrxHCcBojKKwX1terBiRUaqAsFqJiF615XL43r0=
-google.golang.org/genproto v0.0.0-20220228195345-15d65a4533f7 h1:ntPPoHzFW6Xp09ueznmahONZufyoSakK/piXnr2BU3I=
-google.golang.org/genproto v0.0.0-20220228195345-15d65a4533f7/go.mod h1:kGP+zUP2Ddo0ayMi4YuN7C3WZyJvGLZRh8Z5wnAqvEI=
+google.golang.org/genproto v0.0.0-20220422154200-b37d22cd5731 h1:nquqdM9+ps0JZcIiI70+tqoaIFS5Ql4ZuK8UXnz3HfE=
+google.golang.org/genproto v0.0.0-20220422154200-b37d22cd5731/go.mod h1:8w6bsBMX6yCPbAVTeqQHvzxW0EIFigd5lZyahWgyfDo=
 google.golang.org/grpc v1.19.0/go.mod h1:mqu4LbDTu4XGKhr4mRzUsmM4RtVoemTSY81AxZiDr8c=
 google.golang.org/grpc v1.20.1/go.mod h1:10oTOabMzJvdu6/UiuZezV6QK5dSlG84ov/aaiqXj38=
 google.golang.org/grpc v1.21.1/go.mod h1:oYelfM1adQP15Ek0mdvEgi9Df8B9CZIaU1084ijfRaM=
@@ -742,8 +763,9 @@ google.golang.org/grpc v1.33.1/go.mod h1:fr5YgcSWrqhRRxogOsw7RzIpsmvOZ6IcH4kBYTp
 google.golang.org/grpc v1.36.0/go.mod h1:qjiiYl8FncCW8feJPdyg3v6XW24KsRHe+dy9BAGRRjU=
 google.golang.org/grpc v1.38.0/go.mod h1:NREThFqKR1f3iQ6oBuvc5LadQuXVGo9rkm5ZGrQdJfM=
 google.golang.org/grpc v1.41.0/go.mod h1:U3l9uK9J0sini8mHphKoXyaqDA/8VyGnDee1zzIUK6k=
-google.golang.org/grpc v1.44.0 h1:weqSxi/TMs1SqFRMHCtBgXRs8k3X39QIDEZ0pRcttUg=
-google.golang.org/grpc v1.44.0/go.mod h1:k+4IHHFw41K8+bbowsex27ge2rCb65oeWqe4jJ590SU=
+google.golang.org/grpc v1.45.0/go.mod h1:lN7owxKUQEqMfSyQikvvk5tf/6zMPsrK+ONuO11+0rQ=
+google.golang.org/grpc v1.46.0 h1:oCjezcn6g6A75TGoKYBPgKmVBLexhYLM6MebdrPApP8=
+google.golang.org/grpc v1.46.0/go.mod h1:vN9eftEi1UMyUsIF80+uQXhHjbXYbm0uXoFCACuMGWk=
 google.golang.org/protobuf v0.0.0-20200109180630-ec00e32a8dfd/go.mod h1:DFci5gLYBciE7Vtevhsrf46CRTquxDuWsQurQQe4oz8=
 google.golang.org/protobuf v0.0.0-20200221191635-4d8936d0db64/go.mod h1:kwYJMbMJ01Woi6D6+Kah6886xMZcty6N08ah7+eCXa0=
 google.golang.org/protobuf v0.0.0-20200228230310-ab0ca4ff8a60/go.mod h1:cfTl7dwQJ+fmap5saPgwCLgHXTUD7jkjRqWcaiX5VyM=
@@ -756,8 +778,9 @@ google.golang.org/protobuf v1.24.0/go.mod h1:r/3tXBNzIEhYS9I1OUVjXDlt8tc493IdKGj
 google.golang.org/protobuf v1.25.0/go.mod h1:9JNX74DMeImyA3h4bdi1ymwjUzf21/xIlbajtzgsN7c=
 google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
 google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
-google.golang.org/protobuf v1.27.1 h1:SnqbnDw1V7RiZcXPx5MEeqPv2s79L9i7BJUlG/+RurQ=
 google.golang.org/protobuf v1.27.1/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
+google.golang.org/protobuf v1.28.0 h1:w43yiav+6bVFTBQFZX0r7ipe9JQ1QsbMgHwbBziscLw=
+google.golang.org/protobuf v1.28.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
 gopkg.in/alecthomas/kingpin.v2 v2.2.6/go.mod h1:FMv+mEhP44yOT+4EoQTLFTRgOQ1FBLkstjWtayDeSgw=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=

readme-cn.md

@@ -19,7 +19,7 @@
 > 
 > `GOPROXY=https://goproxy.cn/,direct go install github.com/zeromicro/go-zero/tools/goctl@latest`
 > 
-> `goctl migrate —verbose —version v1.3.1`
+> `goctl migrate —verbose —version v1.3.2`
 
 ## 0. go-zero 介绍
 
@@ -116,6 +116,16 @@ GO111MODULE=on GOPROXY=https://goproxy.cn/,direct go get -u github.com/zeromicro
     
     # Go 1.16 及以后版本
     GOPROXY=https://goproxy.cn/,direct go install github.com/zeromicro/go-zero/tools/goctl@latest
+
+    # docker for amd64 architecture
+    docker pull kevinwan/goctl
+    # run goctl like
+    docker run --rm -it -v `pwd`:/app kevinwan/goctl goctl --help
+
+    # docker for arm64 (M1) architecture
+    docker pull kevinwan/goctl:latest-arm64
+    # run goctl like
+    docker run --rm -it -v `pwd`:/app kevinwan/goctl:latest-arm64 goctl --help
     ```
 
     确保 goctl 可执行
@@ -252,6 +262,7 @@ go-zero 已被许多公司用于生产部署，接入场景如在线教育、电
 >58. 七牛云
 >59. 费芮网络
 >60. 51CTO
+>61. 聿旌科技
 
 如果贵公司也已使用 go-zero，欢迎在 [登记地址](https://github.com/zeromicro/go-zero/issues/602) 登记，仅仅为了推广，不做其它用途。
 

readme.md

@@ -18,7 +18,7 @@ English | [简体中文](readme-cn.md)
 > 
 > `go install github.com/zeromicro/go-zero/tools/goctl@latest`
 > 
-> `goctl migrate —verbose —version v1.3.1`
+> `goctl migrate —verbose —version v1.3.2`
 
 ## 0. what is go-zero
 
@@ -119,6 +119,16 @@ go get -u github.com/zeromicro/go-zero
    
    # for Go 1.16 and later
    go install github.com/zeromicro/go-zero/tools/goctl@latest
+
+   # docker for amd64 architecture
+   docker pull kevinwan/goctl
+   # run goctl like
+   docker run --rm -it -v `pwd`:/app kevinwan/goctl goctl --help
+
+   # docker for arm64 (M1) architecture
+   docker pull kevinwan/goctl:latest-arm64
+   # run goctl like
+   docker run --rm -it -v `pwd`:/app kevinwan/goctl:latest-arm64 goctl --help
    ```
 
    make sure goctl is executable.

rest/engine.go

@@ -94,7 +94,7 @@ func (ng *engine) bindRoute(fr featuredRoutes, router httpx.Router, metrics *sta
 		handler.TimeoutHandler(ng.checkedTimeout(fr.timeout)),
 		handler.RecoverHandler,
 		handler.MetricHandler(metrics),
-		handler.MaxBytesHandler(ng.conf.MaxBytes),
+		handler.MaxBytesHandler(ng.checkedMaxBytes(fr.maxBytes)),
 		handler.GunzipHandler,
 	)
 	chain = ng.appendAuthHandler(fr, chain, verifier)
@@ -119,6 +119,14 @@ func (ng *engine) bindRoutes(router httpx.Router) error {
 	return nil
 }
 
+func (ng *engine) checkedMaxBytes(bytes int64) int64 {
+	if bytes > 0 {
+		return bytes
+	}
+
+	return ng.conf.MaxBytes
+}
+
 func (ng *engine) checkedTimeout(timeout time.Duration) time.Duration {
 	if timeout > 0 {
 		return timeout

rest/engine_test.go

@@ -194,6 +194,41 @@ func TestEngine_checkedTimeout(t *testing.T) {
 	}
 }
 
+func TestEngine_checkedMaxBytes(t *testing.T) {
+	tests := []struct {
+		name     string
+		maxBytes int64
+		expect   int64
+	}{
+		{
+			name:   "not set",
+			expect: 1000,
+		},
+		{
+			name:     "less",
+			maxBytes: 500,
+			expect:   500,
+		},
+		{
+			name:     "equal",
+			maxBytes: 1000,
+			expect:   1000,
+		},
+		{
+			name:     "more",
+			maxBytes: 1500,
+			expect:   1500,
+		},
+	}
+
+	ng := newEngine(RestConf{
+		MaxBytes: 1000,
+	})
+	for _, test := range tests {
+		assert.Equal(t, test.expect, ng.checkedMaxBytes(test.maxBytes))
+	}
+}
+
 func TestEngine_notFoundHandler(t *testing.T) {
 	logx.Disable()
 

rest/handler/authhandler.go

@@ -41,7 +41,7 @@ type (
 	AuthorizeOption func(opts *AuthorizeOptions)
 )
 
-// Authorize returns an authorize middleware.
+// Authorize returns an authorization middleware.
 func Authorize(secret string, opts ...AuthorizeOption) func(http.Handler) http.Handler {
 	var authOpts AuthorizeOptions
 	for _, opt := range opts {

rest/handler/cryptionhandler_test.go

@@ -5,6 +5,7 @@ import (
 	"encoding/base64"
 	"io/ioutil"
 	"log"
+	"math/rand"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -116,3 +117,18 @@ func TestCryptionHandler_Hijack(t *testing.T) {
 		writer.Hijack()
 	})
 }
+
+func TestCryptionHandler_ContentTooLong(t *testing.T) {
+	handler := CryptionHandler(aesKey)(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	}))
+	svr := httptest.NewServer(handler)
+	defer svr.Close()
+
+	body := make([]byte, maxBytes+1)
+	rand.Read(body)
+	req, err := http.NewRequest(http.MethodPost, svr.URL, bytes.NewReader(body))
+	assert.Nil(t, err)
+	resp, err := http.DefaultClient.Do(req)
+	assert.Nil(t, err)
+	assert.Equal(t, http.StatusBadRequest, resp.StatusCode)
+}

rest/handler/timeouthandler.go

@@ -20,6 +20,8 @@ import (
 const (
 	statusClientClosedRequest = 499
 	reason                    = "Request Timeout"
+	headerUpgrade             = "Upgrade"
+	valueWebsocket            = "websocket"
 )
 
 // TimeoutHandler returns the handler with given timeout.
@@ -52,6 +54,11 @@ func (h *timeoutHandler) errorBody() string {
 }
 
 func (h *timeoutHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	if r.Header.Get(headerUpgrade) == valueWebsocket {
+		h.handler.ServeHTTP(w, r)
+		return
+	}
+
 	ctx, cancelCtx := context.WithTimeout(r.Context(), h.dt)
 	defer cancelCtx()
 

rest/handler/timeouthandler_test.go

@@ -79,6 +79,19 @@ func TestTimeoutPanic(t *testing.T) {
 	})
 }
 
+func TestTimeoutWebsocket(t *testing.T) {
+	timeoutHandler := TimeoutHandler(time.Millisecond)
+	handler := timeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		time.Sleep(time.Millisecond * 10)
+	}))
+
+	req := httptest.NewRequest(http.MethodGet, "http://localhost", nil)
+	req.Header.Set(headerUpgrade, valueWebsocket)
+	resp := httptest.NewRecorder()
+	handler.ServeHTTP(resp, req)
+	assert.Equal(t, http.StatusOK, resp.Code)
+}
+
 func TestTimeoutWroteHeaderTwice(t *testing.T) {
 	timeoutHandler := TimeoutHandler(time.Minute)
 	handler := timeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -109,7 +122,7 @@ func TestTimeoutWriteBadCode(t *testing.T) {
 func TestTimeoutClientClosed(t *testing.T) {
 	timeoutHandler := TimeoutHandler(time.Minute)
 	handler := timeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		w.WriteHeader(1000)
+		w.WriteHeader(http.StatusServiceUnavailable)
 	}))
 
 	req := httptest.NewRequest(http.MethodGet, "http://localhost", nil)

rest/httpc/internal/interceptor.go

@@ -4,5 +4,5 @@ import "net/http"
 
 type (
 	Interceptor     func(r *http.Request) (*http.Request, ResponseHandler)
-	ResponseHandler func(*http.Response)
+	ResponseHandler func(resp *http.Response, err error)
 )

rest/httpc/internal/loginterceptor.go

@@ -10,15 +10,21 @@ import (
 
 func LogInterceptor(r *http.Request) (*http.Request, ResponseHandler) {
 	start := timex.Now()
-	return r, func(resp *http.Response) {
+	return r, func(resp *http.Response, err error) {
 		duration := timex.Since(start)
+		if err != nil {
+			logger := logx.WithContext(r.Context()).WithDuration(duration)
+			logger.Errorf("[HTTP] %s %s - %v", r.Method, r.URL, err)
+			return
+		}
+
 		var tc propagation.TraceContext
 		ctx := tc.Extract(r.Context(), propagation.HeaderCarrier(resp.Header))
 		logger := logx.WithContext(ctx).WithDuration(duration)
 		if isOkResponse(resp.StatusCode) {
-			logger.Infof("[HTTP] %d - %s %s/%s", resp.StatusCode, r.Method, r.Host, r.RequestURI)
+			logger.Infof("[HTTP] %d - %s %s", resp.StatusCode, r.Method, r.URL)
 		} else {
-			logger.Errorf("[HTTP] %d - %s %s/%s", resp.StatusCode, r.Method, r.Host, r.RequestURI)
+			logger.Errorf("[HTTP] %d - %s %s", resp.StatusCode, r.Method, r.URL)
 		}
 	}
 }

rest/httpc/internal/loginterceptor_test.go

@@ -11,12 +11,13 @@ import (
 func TestLogInterceptor(t *testing.T) {
 	svr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 	}))
+	defer svr.Close()
 	req, err := http.NewRequest(http.MethodGet, svr.URL, nil)
 	assert.Nil(t, err)
 	req, handler := LogInterceptor(req)
 	resp, err := http.DefaultClient.Do(req)
+	handler(resp, err)
 	assert.Nil(t, err)
-	handler(resp)
 	assert.Equal(t, http.StatusOK, resp.StatusCode)
 }
 
@@ -24,11 +25,27 @@ func TestLogInterceptorServerError(t *testing.T) {
 	svr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		w.WriteHeader(http.StatusInternalServerError)
 	}))
+	defer svr.Close()
 	req, err := http.NewRequest(http.MethodGet, svr.URL, nil)
 	assert.Nil(t, err)
 	req, handler := LogInterceptor(req)
 	resp, err := http.DefaultClient.Do(req)
+	handler(resp, err)
 	assert.Nil(t, err)
-	handler(resp)
 	assert.Equal(t, http.StatusInternalServerError, resp.StatusCode)
 }
+
+func TestLogInterceptorServerClosed(t *testing.T) {
+	svr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusInternalServerError)
+	}))
+	defer svr.Close()
+	req, err := http.NewRequest(http.MethodGet, svr.URL, nil)
+	assert.Nil(t, err)
+	svr.Close()
+	req, handler := LogInterceptor(req)
+	resp, err := http.DefaultClient.Do(req)
+	handler(resp, err)
+	assert.NotNil(t, err)
+	assert.Nil(t, resp)
+}

rest/httpc/requests.go

@@ -1,21 +1,166 @@
 package httpc
 
 import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
 	"io"
 	"net/http"
+	nurl "net/url"
+	"strings"
+
+	"github.com/zeromicro/go-zero/core/lang"
+	"github.com/zeromicro/go-zero/core/mapping"
+	"github.com/zeromicro/go-zero/rest/httpc/internal"
+	"github.com/zeromicro/go-zero/rest/internal/header"
 )
 
-// Do sends an HTTP request to the service assocated with the given key.
-func Do(key string, r *http.Request) (*http.Response, error) {
-	return NewService(key).Do(r)
+var interceptors = []internal.Interceptor{
+	internal.LogInterceptor,
 }
 
-// Get sends an HTTP GET request to the service assocated with the given key.
-func Get(key, url string) (*http.Response, error) {
-	return NewService(key).Get(url)
+// Do sends an HTTP request with the given arguments and returns an HTTP response.
+// data is automatically marshal into a *httpRequest, typically it's defined in an API file.
+func Do(ctx context.Context, method, url string, data interface{}) (*http.Response, error) {
+	req, err := buildRequest(ctx, method, url, data)
+	if err != nil {
+		return nil, err
+	}
+
+	return DoRequest(req)
 }
 
-// Post sends an HTTP POST request to the service assocated with the given key.
-func Post(key, url, contentType string, body io.Reader) (*http.Response, error) {
-	return NewService(key).Post(url, contentType, body)
+// DoRequest sends an HTTP request and returns an HTTP response.
+func DoRequest(r *http.Request) (*http.Response, error) {
+	return request(r, defaultClient{})
+}
+
+type (
+	client interface {
+		do(r *http.Request) (*http.Response, error)
+	}
+
+	defaultClient struct{}
+)
+
+func (c defaultClient) do(r *http.Request) (*http.Response, error) {
+	return http.DefaultClient.Do(r)
+}
+
+func buildFormQuery(u *nurl.URL, val map[string]interface{}) string {
+	query := u.Query()
+	for k, v := range val {
+		query.Add(k, fmt.Sprint(v))
+	}
+
+	return query.Encode()
+}
+
+func buildRequest(ctx context.Context, method, url string, data interface{}) (*http.Request, error) {
+	u, err := nurl.Parse(url)
+	if err != nil {
+		return nil, err
+	}
+
+	var val map[string]map[string]interface{}
+	if data != nil {
+		val, err = mapping.Marshal(data)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	if err := fillPath(u, val[pathKey]); err != nil {
+		return nil, err
+	}
+
+	var reader io.Reader
+	jsonVars, hasJsonBody := val[jsonKey]
+	if hasJsonBody {
+		if method == http.MethodGet {
+			return nil, ErrGetWithBody
+		}
+
+		var buf bytes.Buffer
+		enc := json.NewEncoder(&buf)
+		if err := enc.Encode(jsonVars); err != nil {
+			return nil, err
+		}
+
+		reader = &buf
+	}
+
+	req, err := http.NewRequestWithContext(ctx, method, u.String(), reader)
+	if err != nil {
+		return nil, err
+	}
+
+	req.URL.RawQuery = buildFormQuery(u, val[formKey])
+	fillHeader(req, val[headerKey])
+	if hasJsonBody {
+		req.Header.Set(header.ContentType, header.JsonContentType)
+	}
+
+	return req, nil
+}
+
+func fillHeader(r *http.Request, val map[string]interface{}) {
+	for k, v := range val {
+		r.Header.Add(k, fmt.Sprint(v))
+	}
+}
+
+func fillPath(u *nurl.URL, val map[string]interface{}) error {
+	used := make(map[string]lang.PlaceholderType)
+	fields := strings.Split(u.Path, slash)
+
+	for i := range fields {
+		field := fields[i]
+		if len(field) > 0 && field[0] == colon {
+			name := field[1:]
+			ival, ok := val[name]
+			if !ok {
+				return fmt.Errorf("missing path variable %q", name)
+			}
+			value := fmt.Sprint(ival)
+			if len(value) == 0 {
+				return fmt.Errorf("empty path variable %q", name)
+			}
+			fields[i] = value
+			used[name] = lang.Placeholder
+		}
+	}
+
+	if len(val) != len(used) {
+		for key := range used {
+			delete(val, key)
+		}
+
+		var unused []string
+		for key := range val {
+			unused = append(unused, key)
+		}
+
+		return fmt.Errorf("more path variables are provided: %q", strings.Join(unused, ", "))
+	}
+
+	u.Path = strings.Join(fields, slash)
+	return nil
+}
+
+func request(r *http.Request, cli client) (*http.Response, error) {
+	var respHandlers []internal.ResponseHandler
+	for _, interceptor := range interceptors {
+		var h internal.ResponseHandler
+		r, h = interceptor(r)
+		respHandlers = append(respHandlers, h)
+	}
+
+	resp, err := cli.do(r)
+	for i := len(respHandlers) - 1; i >= 0; i-- {
+		respHandlers[i](resp, err)
+	}
+
+	return resp, err
 }

rest/httpc/requests_test.go

@@ -1,37 +1,160 @@
 package httpc
 
 import (
+	"context"
 	"net/http"
 	"net/http/httptest"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/zeromicro/go-zero/rest/httpx"
+	"github.com/zeromicro/go-zero/rest/internal/header"
+	"github.com/zeromicro/go-zero/rest/router"
 )
 
-func TestDo(t *testing.T) {
+func TestDoRequest(t *testing.T) {
 	svr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 	}))
-	_, err := Get("foo", "tcp://bad request")
-	assert.NotNil(t, err)
-	resp, err := Get("foo", svr.URL)
+	defer svr.Close()
+	req, err := http.NewRequest(http.MethodGet, svr.URL, nil)
+	assert.Nil(t, err)
+	resp, err := DoRequest(req)
 	assert.Nil(t, err)
 	assert.Equal(t, http.StatusOK, resp.StatusCode)
 }
 
-func TestDoNotFound(t *testing.T) {
+func TestDoRequest_NotFound(t *testing.T) {
 	svr := httptest.NewServer(http.NotFoundHandler())
-	_, err := Post("foo", "tcp://bad request", "application/json", nil)
-	assert.NotNil(t, err)
-	resp, err := Post("foo", svr.URL, "application/json", nil)
+	defer svr.Close()
+	req, err := http.NewRequest(http.MethodPost, svr.URL, nil)
+	assert.Nil(t, err)
+	req.Header.Set(header.ContentType, header.JsonContentType)
+	resp, err := DoRequest(req)
 	assert.Nil(t, err)
 	assert.Equal(t, http.StatusNotFound, resp.StatusCode)
 }
 
-func TestDoMoved(t *testing.T) {
+func TestDoRequest_Moved(t *testing.T) {
 	svr := httptest.NewServer(http.RedirectHandler("/foo", http.StatusMovedPermanently))
+	defer svr.Close()
 	req, err := http.NewRequest(http.MethodGet, svr.URL, nil)
 	assert.Nil(t, err)
-	_, err = Do("foo", req)
+	_, err = DoRequest(req)
 	// too many redirects
 	assert.NotNil(t, err)
 }
+
+func TestDo(t *testing.T) {
+	type Data struct {
+		Key    string `path:"key"`
+		Value  int    `form:"value"`
+		Header string `header:"X-Header"`
+		Body   string `json:"body"`
+	}
+
+	rt := router.NewRouter()
+	err := rt.Handle(http.MethodPost, "/nodes/:key",
+		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			var req Data
+			assert.Nil(t, httpx.Parse(r, &req))
+		}))
+	assert.Nil(t, err)
+
+	svr := httptest.NewServer(http.HandlerFunc(rt.ServeHTTP))
+	defer svr.Close()
+
+	data := Data{
+		Key:    "foo",
+		Value:  10,
+		Header: "my-header",
+		Body:   "my body",
+	}
+	resp, err := Do(context.Background(), http.MethodPost, svr.URL+"/nodes/:key", data)
+	assert.Nil(t, err)
+	assert.Equal(t, http.StatusOK, resp.StatusCode)
+}
+
+func TestDo_Ptr(t *testing.T) {
+	type Data struct {
+		Key    string `path:"key"`
+		Value  int    `form:"value"`
+		Header string `header:"X-Header"`
+		Body   string `json:"body"`
+	}
+
+	rt := router.NewRouter()
+	err := rt.Handle(http.MethodPost, "/nodes/:key",
+		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			var req Data
+			assert.Nil(t, httpx.Parse(r, &req))
+			assert.Equal(t, "foo", req.Key)
+			assert.Equal(t, 10, req.Value)
+			assert.Equal(t, "my-header", req.Header)
+			assert.Equal(t, "my body", req.Body)
+		}))
+	assert.Nil(t, err)
+
+	svr := httptest.NewServer(http.HandlerFunc(rt.ServeHTTP))
+	defer svr.Close()
+
+	data := &Data{
+		Key:    "foo",
+		Value:  10,
+		Header: "my-header",
+		Body:   "my body",
+	}
+	resp, err := Do(context.Background(), http.MethodPost, svr.URL+"/nodes/:key", data)
+	assert.Nil(t, err)
+	assert.Equal(t, http.StatusOK, resp.StatusCode)
+}
+
+func TestDo_BadRequest(t *testing.T) {
+	_, err := Do(context.Background(), http.MethodPost, ":/nodes/:key", nil)
+	assert.NotNil(t, err)
+
+	val1 := struct {
+		Value string `json:"value,options=[a,b]"`
+	}{
+		Value: "c",
+	}
+	_, err = Do(context.Background(), http.MethodPost, "/nodes/:key", val1)
+	assert.NotNil(t, err)
+
+	val2 := struct {
+		Value string `path:"val"`
+	}{
+		Value: "",
+	}
+	_, err = Do(context.Background(), http.MethodPost, "/nodes/:key", val2)
+	assert.NotNil(t, err)
+
+	val3 := struct {
+		Value string `path:"key"`
+		Body  string `json:"body"`
+	}{
+		Value: "foo",
+	}
+	_, err = Do(context.Background(), http.MethodGet, "/nodes/:key", val3)
+	assert.NotNil(t, err)
+
+	_, err = Do(context.Background(), "\n", "rtmp://nodes", nil)
+	assert.NotNil(t, err)
+
+	val4 := struct {
+		Value string `path:"val"`
+	}{
+		Value: "",
+	}
+	_, err = Do(context.Background(), http.MethodPost, "/nodes/:val", val4)
+	assert.NotNil(t, err)
+
+	val5 := struct {
+		Value   string `path:"val"`
+		Another int    `path:"foo"`
+	}{
+		Value:   "1",
+		Another: 2,
+	}
+	_, err = Do(context.Background(), http.MethodPost, "/nodes/:val", val5)
+	assert.NotNil(t, err)
+}

rest/httpc/responses.go

@@ -0,0 +1,37 @@
+package httpc
+
+import (
+	"net/http"
+	"strings"
+
+	"github.com/zeromicro/go-zero/core/mapping"
+	"github.com/zeromicro/go-zero/rest/internal/encoding"
+	"github.com/zeromicro/go-zero/rest/internal/header"
+)
+
+// Parse parses the response.
+func Parse(resp *http.Response, val interface{}) error {
+	if err := ParseHeaders(resp, val); err != nil {
+		return err
+	}
+
+	return ParseJsonBody(resp, val)
+}
+
+// ParseHeaders parses the rsponse headers.
+func ParseHeaders(resp *http.Response, val interface{}) error {
+	return encoding.ParseHeaders(resp.Header, val)
+}
+
+// ParseJsonBody parses the rsponse body, which should be in json content type.
+func ParseJsonBody(resp *http.Response, val interface{}) error {
+	if withJsonBody(resp) {
+		return mapping.UnmarshalJsonReader(resp.Body, val)
+	}
+
+	return mapping.UnmarshalJsonMap(nil, val)
+}
+
+func withJsonBody(r *http.Response) bool {
+	return r.ContentLength > 0 && strings.Contains(r.Header.Get(header.ContentType), header.ApplicationJson)
+}

rest/httpc/responses_test.go

@@ -0,0 +1,65 @@
+package httpc
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/zeromicro/go-zero/rest/internal/header"
+)
+
+func TestParse(t *testing.T) {
+	var val struct {
+		Foo   string `header:"foo"`
+		Name  string `json:"name"`
+		Value int    `json:"value"`
+	}
+	svr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("foo", "bar")
+		w.Header().Set(header.ContentType, header.JsonContentType)
+		w.Write([]byte(`{"name":"kevin","value":100}`))
+	}))
+	defer svr.Close()
+	req, err := http.NewRequest(http.MethodGet, svr.URL, nil)
+	assert.Nil(t, err)
+	resp, err := DoRequest(req)
+	assert.Nil(t, err)
+	assert.Nil(t, Parse(resp, &val))
+	assert.Equal(t, "bar", val.Foo)
+	assert.Equal(t, "kevin", val.Name)
+	assert.Equal(t, 100, val.Value)
+}
+
+func TestParseHeaderError(t *testing.T) {
+	var val struct {
+		Foo int `header:"foo"`
+	}
+	svr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("foo", "bar")
+		w.Header().Set(header.ContentType, header.JsonContentType)
+	}))
+	defer svr.Close()
+	req, err := http.NewRequest(http.MethodGet, svr.URL, nil)
+	assert.Nil(t, err)
+	resp, err := DoRequest(req)
+	assert.Nil(t, err)
+	assert.NotNil(t, Parse(resp, &val))
+}
+
+func TestParseNoBody(t *testing.T) {
+	var val struct {
+		Foo string `header:"foo"`
+	}
+	svr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("foo", "bar")
+		w.Header().Set(header.ContentType, header.JsonContentType)
+	}))
+	defer svr.Close()
+	req, err := http.NewRequest(http.MethodGet, svr.URL, nil)
+	assert.Nil(t, err)
+	resp, err := DoRequest(req)
+	assert.Nil(t, err)
+	assert.Nil(t, Parse(resp, &val))
+	assert.Equal(t, "bar", val.Foo)
+}

rest/httpc/service.go

@@ -1,33 +1,22 @@
 package httpc
 
 import (
-	"io"
+	"context"
 	"net/http"
 
 	"github.com/zeromicro/go-zero/core/breaker"
-	"github.com/zeromicro/go-zero/core/logx"
-	"github.com/zeromicro/go-zero/rest/httpc/internal"
 )
 
-// ContentType means Content-Type.
-const ContentType = "Content-Type"
-
-var interceptors = []internal.Interceptor{
-	internal.LogInterceptor,
-}
-
 type (
 	// Option is used to customize the *http.Client.
 	Option func(r *http.Request) *http.Request
 
 	// Service represents a remote HTTP service.
 	Service interface {
-		// Do sends an HTTP request to the service.
-		Do(r *http.Request) (*http.Response, error)
-		// Get sends an HTTP GET request to the service.
-		Get(url string) (*http.Response, error)
-		// Post sends an HTTP POST request to the service.
-		Post(url, contentType string, body io.Reader) (*http.Response, error)
+		// Do sends an HTTP request with the given arguments and returns an HTTP response.
+		Do(ctx context.Context, method, url string, data interface{}) (*http.Response, error)
+		// DoRequest sends a HTTP request to the service.
+		DoRequest(r *http.Request) (*http.Response, error)
 	}
 
 	namedService struct {
@@ -53,50 +42,22 @@ func NewServiceWithClient(name string, cli *http.Client, opts ...Option) Service
 	}
 }
 
-// Do sends an HTTP request to the service.
-func (s namedService) Do(r *http.Request) (resp *http.Response, err error) {
-	var respHandlers []internal.ResponseHandler
-	for _, interceptor := range interceptors {
-		var h internal.ResponseHandler
-		r, h = interceptor(r)
-		respHandlers = append(respHandlers, h)
-	}
-
-	resp, err = s.doRequest(r)
-	if err != nil {
-		logx.Errorf("[HTTP] %s %s/%s - %v", r.Method, r.Host, r.RequestURI, err)
-		return
-	}
-
-	for i := len(respHandlers) - 1; i >= 0; i-- {
-		respHandlers[i](resp)
-	}
-
-	return
-}
-
-// Get sends an HTTP GET request to the service.
-func (s namedService) Get(url string) (*http.Response, error) {
-	r, err := http.NewRequest(http.MethodGet, url, nil)
+// Do sends an HTTP request with the given arguments and returns an HTTP response.
+func (s namedService) Do(ctx context.Context, method, url string, data interface{}) (*http.Response, error) {
+	req, err := buildRequest(ctx, method, url, data)
 	if err != nil {
 		return nil, err
 	}
 
-	return s.Do(r)
+	return s.DoRequest(req)
 }
 
-// Post sends an HTTP POST request to the service.
-func (s namedService) Post(url, contentType string, body io.Reader) (*http.Response, error) {
-	r, err := http.NewRequest(http.MethodPost, url, body)
-	if err != nil {
-		return nil, err
-	}
-
-	r.Header.Set(ContentType, contentType)
-	return s.Do(r)
+// DoRequest sends an HTTP request to the service.
+func (s namedService) DoRequest(r *http.Request) (*http.Response, error) {
+	return request(r, s)
 }
 
-func (s namedService) doRequest(r *http.Request) (resp *http.Response, err error) {
+func (s namedService) do(r *http.Request) (resp *http.Response, err error) {
 	for _, opt := range s.opts {
 		r = opt(r)
 	}

rest/httpc/service_test.go

@@ -1,43 +1,86 @@
 package httpc
 
 import (
+	"context"
 	"net/http"
 	"net/http/httptest"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/zeromicro/go-zero/rest/internal/header"
 )
 
-func TestNamedService_Do(t *testing.T) {
+func TestNamedService_DoRequest(t *testing.T) {
 	svr := httptest.NewServer(http.RedirectHandler("/foo", http.StatusMovedPermanently))
+	defer svr.Close()
 	req, err := http.NewRequest(http.MethodGet, svr.URL, nil)
 	assert.Nil(t, err)
 	service := NewService("foo")
-	_, err = service.Do(req)
+	_, err = service.DoRequest(req)
 	// too many redirects
 	assert.NotNil(t, err)
 }
 
-func TestNamedService_Get(t *testing.T) {
+func TestNamedService_DoRequestGet(t *testing.T) {
 	svr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		w.Header().Set("foo", r.Header.Get("foo"))
 	}))
+	defer svr.Close()
 	service := NewService("foo", func(r *http.Request) *http.Request {
 		r.Header.Set("foo", "bar")
 		return r
 	})
-	resp, err := service.Get(svr.URL)
+	req, err := http.NewRequest(http.MethodGet, svr.URL, nil)
+	assert.Nil(t, err)
+	resp, err := service.DoRequest(req)
 	assert.Nil(t, err)
 	assert.Equal(t, http.StatusOK, resp.StatusCode)
 	assert.Equal(t, "bar", resp.Header.Get("foo"))
 }
 
-func TestNamedService_Post(t *testing.T) {
+func TestNamedService_DoRequestPost(t *testing.T) {
 	svr := httptest.NewServer(http.NotFoundHandler())
+	defer svr.Close()
 	service := NewService("foo")
-	_, err := service.Post("tcp://bad request", "application/json", nil)
-	assert.NotNil(t, err)
-	resp, err := service.Post(svr.URL, "application/json", nil)
+	req, err := http.NewRequest(http.MethodPost, svr.URL, nil)
+	assert.Nil(t, err)
+	req.Header.Set(header.ContentType, header.JsonContentType)
+	resp, err := service.DoRequest(req)
 	assert.Nil(t, err)
 	assert.Equal(t, http.StatusNotFound, resp.StatusCode)
 }
+
+func TestNamedService_Do(t *testing.T) {
+	type Data struct {
+		Key    string `path:"key"`
+		Value  int    `form:"value"`
+		Header string `header:"X-Header"`
+		Body   string `json:"body"`
+	}
+
+	svr := httptest.NewServer(http.NotFoundHandler())
+	defer svr.Close()
+
+	service := NewService("foo")
+	data := Data{
+		Key:    "foo",
+		Value:  10,
+		Header: "my-header",
+		Body:   "my body",
+	}
+	resp, err := service.Do(context.Background(), http.MethodPost, svr.URL+"/nodes/:key", data)
+	assert.Nil(t, err)
+	assert.Equal(t, http.StatusNotFound, resp.StatusCode)
+}
+
+func TestNamedService_DoBadRequest(t *testing.T) {
+	val := struct {
+		Value string `json:"value,options=[a,b]"`
+	}{
+		Value: "c",
+	}
+
+	service := NewService("foo")
+	_, err := service.Do(context.Background(), http.MethodPost, "/nodes/:key", val)
+	assert.NotNil(t, err)
+}

rest/httpc/vars.go

@@ -0,0 +1,15 @@
+package httpc
+
+import "errors"
+
+const (
+	pathKey   = "path"
+	formKey   = "form"
+	headerKey = "header"
+	jsonKey   = "json"
+	slash     = "/"
+	colon     = ':'
+)
+
+// ErrGetWithBody indicates that GET request with body.
+var ErrGetWithBody = errors.New("HTTP GET should not have body")

rest/httpx/requests.go

@@ -3,17 +3,17 @@ package httpx
 import (
 	"io"
 	"net/http"
-	"net/textproto"
 	"strings"
 
 	"github.com/zeromicro/go-zero/core/mapping"
+	"github.com/zeromicro/go-zero/rest/internal/encoding"
+	"github.com/zeromicro/go-zero/rest/internal/header"
 	"github.com/zeromicro/go-zero/rest/pathvar"
 )
 
 const (
 	formKey           = "form"
 	pathKey           = "path"
-	headerKey         = "header"
 	maxMemory         = 32 << 20 // 32MB
 	maxBodyLen        = 8 << 20  // 8MB
 	separator         = ";"
@@ -21,10 +21,8 @@ const (
 )
 
 var (
-	formUnmarshaler   = mapping.NewUnmarshaler(formKey, mapping.WithStringValues())
-	pathUnmarshaler   = mapping.NewUnmarshaler(pathKey, mapping.WithStringValues())
-	headerUnmarshaler = mapping.NewUnmarshaler(headerKey, mapping.WithStringValues(),
-		mapping.WithCanonicalKeyFunc(textproto.CanonicalMIMEHeaderKey))
+	formUnmarshaler = mapping.NewUnmarshaler(formKey, mapping.WithStringValues())
+	pathUnmarshaler = mapping.NewUnmarshaler(pathKey, mapping.WithStringValues())
 )
 
 // Parse parses the request.
@@ -46,16 +44,7 @@ func Parse(r *http.Request, v interface{}) error {
 
 // ParseHeaders parses the headers request.
 func ParseHeaders(r *http.Request, v interface{}) error {
-	m := map[string]interface{}{}
-	for k, v := range r.Header {
-		if len(v) == 1 {
-			m[k] = v[0]
-		} else {
-			m[k] = v
-		}
-	}
-
-	return headerUnmarshaler.Unmarshal(m, v)
+	return encoding.ParseHeaders(r.Header, v)
 }
 
 // ParseForm parses the form request.
@@ -126,5 +115,5 @@ func ParsePath(r *http.Request, v interface{}) error {
 }
 
 func withJsonBody(r *http.Request) bool {
-	return r.ContentLength > 0 && strings.Contains(r.Header.Get(ContentType), ApplicationJson)
+	return r.ContentLength > 0 && strings.Contains(r.Header.Get(header.ContentType), header.ApplicationJson)
 }

rest/httpx/requests_test.go

@@ -8,6 +8,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/zeromicro/go-zero/rest/internal/header"
 	"github.com/zeromicro/go-zero/rest/pathvar"
 )
 
@@ -204,7 +205,7 @@ func TestParseJsonBody(t *testing.T) {
 
 		body := `{"name":"kevin", "age": 18}`
 		r := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body))
-		r.Header.Set(ContentType, ApplicationJson)
+		r.Header.Set(ContentType, header.JsonContentType)
 
 		assert.Nil(t, Parse(r, &v))
 		assert.Equal(t, "kevin", v.Name)
@@ -267,7 +268,7 @@ func TestParseHeaders(t *testing.T) {
 	request.Header.Add("addrs", "addr2")
 	request.Header.Add("X-Forwarded-For", "10.0.10.11")
 	request.Header.Add("x-real-ip", "10.0.11.10")
-	request.Header.Add("Accept", "application/json")
+	request.Header.Add("Accept", header.JsonContentType)
 	err = ParseHeaders(request, &v)
 	if err != nil {
 		t.Fatal(err)
@@ -277,7 +278,7 @@ func TestParseHeaders(t *testing.T) {
 	assert.Equal(t, []string{"addr1", "addr2"}, v.Addrs)
 	assert.Equal(t, "10.0.10.11", v.XForwardedFor)
 	assert.Equal(t, "10.0.11.10", v.XRealIP)
-	assert.Equal(t, "application/json", v.Accept)
+	assert.Equal(t, header.JsonContentType, v.Accept)
 }
 
 func TestParseHeaders_Error(t *testing.T) {

rest/httpx/responses.go

@@ -6,6 +6,7 @@ import (
 	"sync"
 
 	"github.com/zeromicro/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/rest/internal/header"
 )
 
 var (
@@ -61,12 +62,16 @@ func SetErrorHandler(handler func(error) (int, interface{})) {
 
 // WriteJson writes v as json string into w with code.
 func WriteJson(w http.ResponseWriter, code int, v interface{}) {
-	w.Header().Set(ContentType, ApplicationJson)
+	bs, err := json.Marshal(v)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	w.Header().Set(ContentType, header.JsonContentType)
 	w.WriteHeader(code)
 
-	if bs, err := json.Marshal(v); err != nil {
-		http.Error(w, err.Error(), http.StatusInternalServerError)
-	} else if n, err := w.Write(bs); err != nil {
+	if n, err := w.Write(bs); err != nil {
 		// http.ErrHandlerTimeout has been handled by http.TimeoutHandler,
 		// so it's ignored here.
 		if err != http.ErrHandlerTimeout {

rest/httpx/responses_test.go

@@ -146,6 +146,16 @@ func TestWriteJsonLessWritten(t *testing.T) {
 	assert.Equal(t, http.StatusOK, w.code)
 }
 
+func TestWriteJsonMarshalFailed(t *testing.T) {
+	w := tracedResponseWriter{
+		headers: make(map[string][]string),
+	}
+	WriteJson(&w, http.StatusOK, map[string]interface{}{
+		"Data": complex(0, 0),
+	})
+	assert.Equal(t, http.StatusInternalServerError, w.code)
+}
+
 type tracedResponseWriter struct {
 	headers     map[string][]string
 	builder     strings.Builder
@@ -153,6 +163,7 @@ type tracedResponseWriter struct {
 	code        int
 	lessWritten bool
 	timeout     bool
+	wroteHeader bool
 }
 
 func (w *tracedResponseWriter) Header() http.Header {
@@ -174,5 +185,9 @@ func (w *tracedResponseWriter) Write(bytes []byte) (n int, err error) {
 }
 
 func (w *tracedResponseWriter) WriteHeader(code int) {
+	if w.wroteHeader {
+		return
+	}
+	w.wroteHeader = true
 	w.code = code
 }

rest/httpx/vars.go

@@ -1,14 +1,16 @@
 package httpx
 
+import "github.com/zeromicro/go-zero/rest/internal/header"
+
 const (
-	// ApplicationJson means application/json.
-	ApplicationJson = "application/json"
 	// ContentEncoding means Content-Encoding.
 	ContentEncoding = "Content-Encoding"
 	// ContentSecurity means X-Content-Security.
 	ContentSecurity = "X-Content-Security"
 	// ContentType means Content-Type.
-	ContentType = "Content-Type"
+	ContentType = header.ContentType
+	// JsonContentType means application/json.
+	JsonContentType = header.JsonContentType
 	// KeyField means key.
 	KeyField = "key"
 	// SecretField means secret.

rest/internal/cors/handlers.go

@@ -2,6 +2,7 @@ package cors
 
 import (
 	"net/http"
+	"strings"
 
 	"github.com/zeromicro/go-zero/rest/internal/response"
 )
@@ -81,7 +82,7 @@ func isOriginAllowed(allows []string, origin string) bool {
 			return true
 		}
 
-		if o == origin {
+		if strings.HasSuffix(origin, o) {
 			return true
 		}
 	}

rest/internal/cors/handlers_test.go

@@ -31,6 +31,12 @@ func TestCorsHandlerWithOrigins(t *testing.T) {
 			reqOrigin: "http://local",
 			expect:    "http://local",
 		},
+		{
+			name:      "allow sub origins",
+			origins:   []string{"local", "remote"},
+			reqOrigin: "sub.local",
+			expect:    "sub.local",
+		},
 		{
 			name:      "allow all origins",
 			reqOrigin: "http://local",

rest/internal/encoding/parser.go

@@ -0,0 +1,27 @@
+package encoding
+
+import (
+	"net/http"
+	"net/textproto"
+
+	"github.com/zeromicro/go-zero/core/mapping"
+)
+
+const headerKey = "header"
+
+var headerUnmarshaler = mapping.NewUnmarshaler(headerKey, mapping.WithStringValues(),
+	mapping.WithCanonicalKeyFunc(textproto.CanonicalMIMEHeaderKey))
+
+// ParseHeaders parses the headers request.
+func ParseHeaders(header http.Header, v interface{}) error {
+	m := map[string]interface{}{}
+	for k, v := range header {
+		if len(v) == 1 {
+			m[k] = v[0]
+		} else {
+			m[k] = v
+		}
+	}
+
+	return headerUnmarshaler.Unmarshal(m, v)
+}

rest/internal/encoding/parser_test.go

@@ -0,0 +1,40 @@
+package encoding
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestParseHeaders(t *testing.T) {
+	var val struct {
+		Foo string `header:"foo"`
+		Baz int    `header:"baz"`
+		Qux bool   `header:"qux,default=true"`
+	}
+	r := httptest.NewRequest(http.MethodGet, "/any", nil)
+	r.Header.Set("foo", "bar")
+	r.Header.Set("baz", "1")
+	assert.Nil(t, ParseHeaders(r.Header, &val))
+	assert.Equal(t, "bar", val.Foo)
+	assert.Equal(t, 1, val.Baz)
+	assert.True(t, val.Qux)
+}
+
+func TestParseHeadersMulti(t *testing.T) {
+	var val struct {
+		Foo []string `header:"foo"`
+		Baz int      `header:"baz"`
+		Qux bool     `header:"qux,default=true"`
+	}
+	r := httptest.NewRequest(http.MethodGet, "/any", nil)
+	r.Header.Set("foo", "bar")
+	r.Header.Add("foo", "bar1")
+	r.Header.Set("baz", "1")
+	assert.Nil(t, ParseHeaders(r.Header, &val))
+	assert.Equal(t, []string{"bar", "bar1"}, val.Foo)
+	assert.Equal(t, 1, val.Baz)
+	assert.True(t, val.Qux)
+}

rest/internal/header/headers.go

@@ -0,0 +1,10 @@
+package header
+
+const (
+	// ApplicationJson stands for application/json.
+	ApplicationJson = "application/json"
+	// ContentType is the header key for Content-Type.
+	ContentType = "Content-Type"
+	// JsonContentType is the content type for JSON.
+	JsonContentType = "application/json; charset=utf-8"
+)

rest/router/patrouter_test.go

@@ -11,13 +11,11 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/zeromicro/go-zero/rest/httpx"
+	"github.com/zeromicro/go-zero/rest/internal/header"
 	"github.com/zeromicro/go-zero/rest/pathvar"
 )
 
-const (
-	applicationJsonWithUtf8 = "application/json; charset=utf-8"
-	contentLength           = "Content-Length"
-)
+const contentLength = "Content-Length"
 
 type mockedResponseWriter struct {
 	code int
@@ -167,7 +165,7 @@ func TestParseJsonPost(t *testing.T) {
 	r, err := http.NewRequest(http.MethodPost, "http://hello.com/kevin/2017?nickname=whatever&zipcode=200000",
 		bytes.NewBufferString(`{"location": "shanghai", "time": 20170912}`))
 	assert.Nil(t, err)
-	r.Header.Set(httpx.ContentType, httpx.ApplicationJson)
+	r.Header.Set(httpx.ContentType, httpx.JsonContentType)
 
 	router := NewRouter()
 	err = router.Handle(http.MethodPost, "/:name/:year", http.HandlerFunc(func(
@@ -199,7 +197,7 @@ func TestParseJsonPostWithIntSlice(t *testing.T) {
 	r, err := http.NewRequest(http.MethodPost, "http://hello.com/kevin/2017",
 		bytes.NewBufferString(`{"ages": [1, 2], "years": [3, 4]}`))
 	assert.Nil(t, err)
-	r.Header.Set(httpx.ContentType, httpx.ApplicationJson)
+	r.Header.Set(httpx.ContentType, httpx.JsonContentType)
 
 	router := NewRouter()
 	err = router.Handle(http.MethodPost, "/:name/:year", http.HandlerFunc(func(
@@ -227,7 +225,7 @@ func TestParseJsonPostError(t *testing.T) {
 	r, err := http.NewRequest(http.MethodPost, "http://hello.com/kevin/2017?nickname=whatever&zipcode=200000",
 		bytes.NewBufferString(payload))
 	assert.Nil(t, err)
-	r.Header.Set(httpx.ContentType, httpx.ApplicationJson)
+	r.Header.Set(httpx.ContentType, httpx.JsonContentType)
 
 	router := NewRouter()
 	err = router.Handle(http.MethodPost, "/:name/:year", http.HandlerFunc(
@@ -255,7 +253,7 @@ func TestParseJsonPostInvalidRequest(t *testing.T) {
 	r, err := http.NewRequest(http.MethodPost, "http://hello.com/",
 		bytes.NewBufferString(payload))
 	assert.Nil(t, err)
-	r.Header.Set(httpx.ContentType, httpx.ApplicationJson)
+	r.Header.Set(httpx.ContentType, httpx.JsonContentType)
 
 	router := NewRouter()
 	err = router.Handle(http.MethodPost, "/", http.HandlerFunc(
@@ -277,7 +275,7 @@ func TestParseJsonPostRequired(t *testing.T) {
 	r, err := http.NewRequest(http.MethodPost, "http://hello.com/kevin/2017",
 		bytes.NewBufferString(`{"location": "shanghai"`))
 	assert.Nil(t, err)
-	r.Header.Set(httpx.ContentType, httpx.ApplicationJson)
+	r.Header.Set(httpx.ContentType, httpx.JsonContentType)
 
 	router := NewRouter()
 	err = router.Handle(http.MethodPost, "/:name/:year", http.HandlerFunc(
@@ -455,7 +453,7 @@ func TestParsePtrInRequest(t *testing.T) {
 	r, err := http.NewRequest(http.MethodPost, "http://hello.com/kevin/2017",
 		bytes.NewBufferString(`{"audio": {"volume": 100}}`))
 	assert.Nil(t, err)
-	r.Header.Set(httpx.ContentType, httpx.ApplicationJson)
+	r.Header.Set(httpx.ContentType, httpx.JsonContentType)
 
 	type (
 		Request struct {
@@ -603,7 +601,7 @@ func TestParseWrappedRequest(t *testing.T) {
 func TestParseWrappedGetRequestWithJsonHeader(t *testing.T) {
 	r, err := http.NewRequest(http.MethodGet, "http://hello.com/kevin/2017", nil)
 	assert.Nil(t, err)
-	r.Header.Set(httpx.ContentType, applicationJsonWithUtf8)
+	r.Header.Set(httpx.ContentType, header.JsonContentType)
 
 	type (
 		Request struct {
@@ -636,7 +634,7 @@ func TestParseWrappedGetRequestWithJsonHeader(t *testing.T) {
 func TestParseWrappedHeadRequestWithJsonHeader(t *testing.T) {
 	r, err := http.NewRequest(http.MethodHead, "http://hello.com/kevin/2017", nil)
 	assert.Nil(t, err)
-	r.Header.Set(httpx.ContentType, applicationJsonWithUtf8)
+	r.Header.Set(httpx.ContentType, header.JsonContentType)
 
 	type (
 		Request struct {
@@ -702,7 +700,7 @@ func TestParseWithAll(t *testing.T) {
 	r, err := http.NewRequest(http.MethodPost, "http://hello.com/kevin/2017?nickname=whatever&zipcode=200000",
 		bytes.NewBufferString(`{"location": "shanghai", "time": 20170912}`))
 	assert.Nil(t, err)
-	r.Header.Set(httpx.ContentType, httpx.ApplicationJson)
+	r.Header.Set(httpx.ContentType, httpx.JsonContentType)
 
 	router := NewRouter()
 	err = router.Handle(http.MethodPost, "/:name/:year", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -733,7 +731,7 @@ func TestParseWithAllUtf8(t *testing.T) {
 	r, err := http.NewRequest(http.MethodPost, "http://hello.com/kevin/2017?nickname=whatever&zipcode=200000",
 		bytes.NewBufferString(`{"location": "shanghai", "time": 20170912}`))
 	assert.Nil(t, err)
-	r.Header.Set(httpx.ContentType, applicationJsonWithUtf8)
+	r.Header.Set(httpx.ContentType, header.JsonContentType)
 
 	router := NewRouter()
 	err = router.Handle(http.MethodPost, "/:name/:year", http.HandlerFunc(
@@ -923,7 +921,7 @@ func TestParseWithMissingAllPaths(t *testing.T) {
 func TestParseGetWithContentLengthHeader(t *testing.T) {
 	r, err := http.NewRequest(http.MethodGet, "http://hello.com/kevin/2017?nickname=whatever&zipcode=200000", nil)
 	assert.Nil(t, err)
-	r.Header.Set(httpx.ContentType, httpx.ApplicationJson)
+	r.Header.Set(httpx.ContentType, header.JsonContentType)
 	r.Header.Set(contentLength, "1024")
 
 	router := NewRouter()
@@ -951,7 +949,7 @@ func TestParseJsonPostWithTypeMismatch(t *testing.T) {
 	r, err := http.NewRequest(http.MethodPost, "http://hello.com/kevin/2017?nickname=whatever&zipcode=200000",
 		bytes.NewBufferString(`{"time": "20170912"}`))
 	assert.Nil(t, err)
-	r.Header.Set(httpx.ContentType, applicationJsonWithUtf8)
+	r.Header.Set(httpx.ContentType, header.JsonContentType)
 
 	router := NewRouter()
 	err = router.Handle(http.MethodPost, "/:name/:year", http.HandlerFunc(
@@ -977,7 +975,7 @@ func TestParseJsonPostWithInt2String(t *testing.T) {
 	r, err := http.NewRequest(http.MethodPost, "http://hello.com/kevin/2017",
 		bytes.NewBufferString(`{"time": 20170912}`))
 	assert.Nil(t, err)
-	r.Header.Set(httpx.ContentType, applicationJsonWithUtf8)
+	r.Header.Set(httpx.ContentType, header.JsonContentType)
 
 	router := NewRouter()
 	err = router.Handle(http.MethodPost, "/:name/:year", http.HandlerFunc(

rest/server.go

@@ -137,6 +137,13 @@ func WithJwtTransition(secret, prevSecret string) RouteOption {
 	}
 }
 
+// WithMaxBytes returns a RouteOption to set maxBytes with the given value.
+func WithMaxBytes(maxBytes int64) RouteOption {
+	return func(r *featuredRoutes) {
+		r.maxBytes = maxBytes
+	}
+}
+
 // WithMiddlewares adds given middlewares to given routes.
 func WithMiddlewares(ms []Middleware, rs ...Route) []Route {
 	for i := len(ms) - 1; i >= 0; i-- {

rest/server_test.go

@@ -95,6 +95,13 @@ Port: 54321
 	}
 }
 
+func TestWithMaxBytes(t *testing.T) {
+	const maxBytes = 1000
+	var fr featuredRoutes
+	WithMaxBytes(maxBytes)(&fr)
+	assert.Equal(t, int64(maxBytes), fr.maxBytes)
+}
+
 func TestWithMiddleware(t *testing.T) {
 	m := make(map[string]string)
 	rt := router.NewRouter()

rest/types.go

@@ -36,5 +36,6 @@ type (
 		jwt       jwtSetting
 		signature signatureSetting
 		routes    []Route
+		maxBytes  int64
 	}
 )

tools/goctl/Dockerfile

@@ -0,0 +1,34 @@
+FROM golang:alpine AS builder
+
+LABEL stage=gobuilder
+
+ENV CGO_ENABLED 0
+ENV GOPROXY https://goproxy.cn,direct
+
+RUN apk update --no-cache && apk add --no-cache tzdata
+RUN go install google.golang.org/protobuf/cmd/protoc-gen-go@latest
+RUN go install google.golang.org/grpc/cmd/protoc-gen-go-grpc@latest
+
+WORKDIR /build
+
+ADD go.mod .
+ADD go.sum .
+RUN go mod download
+COPY . .
+RUN go build -ldflags="-s -w" -o /app/goctl ./goctl.go
+
+
+FROM golang:alpine
+
+RUN apk update --no-cache && apk add --no-cache protoc
+
+COPY --from=builder /etc/ssl/certs/ca-certificates.crt /etc/ssl/certs/ca-certificates.crt
+COPY --from=builder /usr/share/zoneinfo/Asia/Shanghai /usr/share/zoneinfo/Asia/Shanghai
+COPY --from=builder /go/bin/protoc-gen-go /usr/bin/protoc-gen-go
+COPY --from=builder /go/bin/protoc-gen-go-grpc /usr/bin/protoc-gen-go-grpc
+ENV TZ Asia/Shanghai
+
+WORKDIR /app
+COPY --from=builder /app/goctl /usr/bin/goctl
+
+CMD ["goctl"]

tools/goctl/Makefile

@@ -1,12 +1,25 @@
 build:
 	go build -ldflags="-s -w" goctl.go
 	$(if $(shell command -v upx), upx goctl)
+
 mac:
 	GOOS=darwin go build -ldflags="-s -w" -o goctl-darwin goctl.go
 	$(if $(shell command -v upx), upx goctl-darwin)
+
 win:
 	GOOS=windows go build -ldflags="-s -w" -o goctl.exe goctl.go
 	$(if $(shell command -v upx), upx goctl.exe)
+
 linux:
 	GOOS=linux go build -ldflags="-s -w" -o goctl-linux goctl.go
 	$(if $(shell command -v upx), upx goctl-linux)
+
+image:
+	docker build --rm --platform linux/amd64 -t kevinwan/goctl:$(version) .
+	docker tag kevinwan/goctl:$(version) kevinwan/goctl:latest
+	docker push kevinwan/goctl:$(version)
+	docker push kevinwan/goctl:latest
+	docker build --rm --platform linux/arm64 -t kevinwan/goctl:$(version)-arm64 .
+	docker tag kevinwan/goctl:$(version)-arm64 kevinwan/goctl:latest-arm64
+	docker push kevinwan/goctl:$(version)-arm64
+	docker push kevinwan/goctl:latest-arm64

tools/goctl/api/apigen/api.tpl

@@ -0,0 +1,24 @@
+syntax = "v1"
+
+info (
+	title: // TODO: add title
+	desc: // TODO: add description
+	author: "{{.gitUser}}"
+	email: "{{.gitEmail}}"
+)
+
+type request {
+	// TODO: add members here and delete this comment
+}
+
+type response {
+	// TODO: add members here and delete this comment
+}
+
+service {{.serviceName}} {
+	@handler GetUser // TODO: set handler name and delete this comment
+	get /users/id/:userId(request) returns(response)
+
+	@handler CreateUser // TODO: set handler name and delete this comment
+	post /users/create(request)
+}

tools/goctl/api/apigen/gen.go

@@ -1,11 +1,12 @@
 package apigen
 
 import (
+	_ "embed"
 	"errors"
 	"fmt"
+	"html/template"
 	"path/filepath"
 	"strings"
-	"text/template"
 
 	"github.com/logrusorgru/aurora"
 	"github.com/urfave/cli"
@@ -13,35 +14,15 @@ import (
 	"github.com/zeromicro/go-zero/tools/goctl/util/pathx"
 )
 
-const apiTemplate = `
-syntax = "v1"
-
-info (
-	title: // TODO: add title
-	desc: // TODO: add description
-	author: "{{.gitUser}}"
-	email: "{{.gitEmail}}"
-)
-
-type request {
-	// TODO: add members here and delete this comment
-}
-
-type response {
-	// TODO: add members here and delete this comment
-}
-
-service {{.serviceName}} {
-	@handler GetUser // TODO: set handler name and delete this comment
-	get /users/id/:userId(request) returns(response)
-
-	@handler CreateUser // TODO: set handler name and delete this comment
-	post /users/create(request)
-}
-`
+//go:embed api.tpl
+var apiTemplate string
 
 // ApiCommand create api template file
 func ApiCommand(c *cli.Context) error {
+	if c.NumFlags() == 0 {
+		cli.ShowAppHelpAndExit(c, 1)
+	}
+
 	apiFile := c.String("o")
 	if len(apiFile) == 0 {
 		return errors.New("missing -o")

tools/goctl/api/dartgen/gen.go

@@ -32,6 +32,10 @@ func DartCommand(c *cli.Context) error {
 		return err
 	}
 
+	if err := api.Validate(); err != nil {
+		return err
+	}
+
 	api.Service = api.Service.JoinPrefix()
 	if !strings.HasSuffix(dir, "/") {
 		dir = dir + "/"

tools/goctl/api/dartgen/vars.go

@@ -58,7 +58,7 @@ Future _apiRequest(String method, String path, dynamic data,
       r = await client.getUrl(Uri.parse('https://' + serverHost + path));
     }
 
-    r.headers.set('Content-Type', 'application/json');
+    r.headers.set('Content-Type', 'application/json; charset=utf-8');
     if (tokens != null) {
       r.headers.set('Authorization', tokens.accessToken);
     }
@@ -149,7 +149,7 @@ Future _apiRequest(String method, String path, dynamic data,
 				r = await client.getUrl(Uri.parse('https://' + serverHost + path));
 			}
 
-			r.headers.set('Content-Type', 'application/json');
+			r.headers.set('Content-Type', 'application/json; charset=utf-8');
 			if (tokens != null) {
 				r.headers.set('Authorization', tokens.accessToken);
 			}

tools/goctl/api/docgen/doc.go

@@ -2,6 +2,7 @@ package docgen
 
 import (
 	"bytes"
+	_ "embed"
 	"fmt"
 	"html/template"
 	"strconv"
@@ -13,25 +14,8 @@ import (
 	"github.com/zeromicro/go-zero/tools/goctl/api/util"
 )
 
-const (
-	markdownTemplate = `
-### {{.index}}. {{.routeComment}}
-
-1. 路由定义
-
-- Url: {{.uri}}
-- Method: {{.method}}
-- Request: {{.requestType}}
-- Response: {{.responseType}}
-
-2. 请求定义
-{{.requestContent}}
-
-3. 返回定义
-{{.responseContent}}  
-
-`
-)
+//go:embed markdown.tpl
+var markdownTemplate string
 
 func genDoc(api *spec.ApiSpec, dir, filename string) error {
 	fp, _, err := util.MaybeCreateFile(dir, "", filename)

tools/goctl/api/docgen/markdown.tpl

@@ -0,0 +1,16 @@
+### {{.index}}. {{.routeComment}}
+
+1. route definition
+
+- Url: {{.uri}}
+- Method: {{.method}}
+- Request: {{.requestType}}
+- Response: {{.responseType}}
+
+2. request definition
+
+{{.requestContent}}
+
+3. response definition
+
+{{.responseContent}}

tools/goctl/api/format/format.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"go/format"
 	"go/scanner"
+	"io"
 	"io/ioutil"
 	"os"
 	"path/filepath"
@@ -29,14 +30,14 @@ const (
 func GoFormatApi(c *cli.Context) error {
 	useStdin := c.Bool("stdin")
 	skipCheckDeclare := c.Bool("declare")
+	dir := c.String("dir")
 
 	var be errorx.BatchError
 	if useStdin {
-		if err := apiFormatByStdin(skipCheckDeclare); err != nil {
+		if err := apiFormatReader(os.Stdin, dir, skipCheckDeclare); err != nil {
 			be.Add(err)
 		}
 	} else {
-		dir := c.String("dir")
 		if len(dir) == 0 {
 			return errors.New("missing -dir")
 		}
@@ -65,13 +66,14 @@ func GoFormatApi(c *cli.Context) error {
 	return be.Err()
 }
 
-func apiFormatByStdin(skipCheckDeclare bool) error {
-	data, err := ioutil.ReadAll(os.Stdin)
+// apiFormatReader
+// filename is needed when there are `import` literals.
+func apiFormatReader(reader io.Reader, filename string, skipCheckDeclare bool) error {
+	data, err := ioutil.ReadAll(reader)
 	if err != nil {
 		return err
 	}
-
-	result, err := apiFormat(string(data), skipCheckDeclare)
+	result, err := apiFormat(string(data), skipCheckDeclare, filename)
 	if err != nil {
 		return err
 	}

tools/goctl/api/format/format_test.go

@@ -1,15 +1,21 @@
 package format
 
 import (
+	"fmt"
+	"io/fs"
+	"io/ioutil"
+	"os"
+	"path"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 const (
 	notFormattedStr = `
 type Request struct {
-  Name string ` + "`" + `path:"name,options=you|me"` + "`" + `  
+  Name string ` + "`" + `path:"name,options=you|me"` + "`" + `
 }
 type Response struct {
   Message string ` + "`" + `json:"message"` + "`" + `
@@ -45,3 +51,26 @@ func TestFormat(t *testing.T) {
 	_, err = apiFormat(notFormattedStr, false)
 	assert.Errorf(t, err, " line 7:13 can not found declaration 'Student' in context")
 }
+
+func Test_apiFormatReader_issue1721(t *testing.T) {
+	dir, err := os.MkdirTemp("", "goctl-api-format")
+	require.NoError(t, err)
+	defer os.RemoveAll(dir)
+	subDir := path.Join(dir, "sub")
+	err = os.MkdirAll(subDir, fs.ModePerm)
+	require.NoError(t, err)
+
+	importedFilename := path.Join(dir, "foo.api")
+	err = ioutil.WriteFile(importedFilename, []byte{}, fs.ModePerm)
+	require.NoError(t, err)
+
+	filename := path.Join(subDir, "bar.api")
+	err = ioutil.WriteFile(filename, []byte(fmt.Sprintf(`import "%s"`, importedFilename)), 0644)
+	require.NoError(t, err)
+
+	f, err := os.Open(filename)
+	require.NoError(t, err)
+
+	err = apiFormatReader(f, filename, false)
+	assert.NoError(t, err)
+}

tools/goctl/api/gogen/etc.tpl

@@ -0,0 +1,3 @@
+Name: {{.serviceName}}
+Host: {{.host}}
+Port: {{.port}}

tools/goctl/api/gogen/gen.go

@@ -61,6 +61,10 @@ func DoGenProject(apiFile, dir, style string) error {
 		return err
 	}
 
+	if err := api.Validate(); err != nil {
+		return err
+	}
+
 	cfg, err := config.NewConfig(style)
 	if err != nil {
 		return err

tools/goctl/api/gogen/gen_test.go

@@ -1,6 +1,7 @@
 package gogen
 
 import (
+	_ "embed"
 	goformat "go/format"
 	"io/ioutil"
 	"os"
@@ -10,314 +11,45 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/zeromicro/go-zero/tools/goctl/api/parser"
+	"github.com/zeromicro/go-zero/tools/goctl/rpc/execx"
+	"github.com/zeromicro/go-zero/tools/goctl/util/pathx"
 )
 
-const testApiTemplate = `
-info(
-    title: doc title
-    desc: ">
-    doc description first part,
-    doc description second part<"
-    version: 1.0
+var (
+	//go:embed testdata/test_api_template.api
+	testApiTemplate string
+	//go:embed testdata/test_multi_service_template.api
+	testMultiServiceTemplate string
+	//go:embed testdata/ap_ino_info.api
+	apiNoInfo string
+	//go:embed testdata/invalid_api_file.api
+	invalidApiFile string
+	//go:embed testdata/anonymous_annotation.api
+	anonymousAnnotation string
+	//go:embed testdata/api_has_middleware.api
+	apiHasMiddleware string
+	//go:embed testdata/api_jwt.api
+	apiJwt string
+	//go:embed testdata/api_jwt_with_middleware.api
+	apiJwtWithMiddleware string
+	//go:embed testdata/api_has_no_request.api
+	apiHasNoRequest string
+	//go:embed testdata/api_route_test.api
+	apiRouteTest string
+	//go:embed testdata/has_comment_api_test.api
+	hasCommentApiTest string
+	//go:embed testdata/has_inline_no_exist_test.api
+	hasInlineNoExistTest string
+	//go:embed testdata/import_api.api
+	importApi string
+	//go:embed testdata/has_import_api.api
+	hasImportApi string
+	//go:embed testdata/no_struct_tag_api.api
+	noStructTagApi string
+	//go:embed testdata/nest_type_api.api
+	nestTypeApi string
 )
 
-// TODO: test
-// {
-type Request struct {  // TODO: test
-  // TODO
-  Name string ` + "`" + `path:"name,options=you|me"` + "`" + `   // }
-} // TODO: test
-
-// TODO: test
-type Response struct {
-  Message string ` + "`" + `json:"message"` + "`" + `
-}
-
-@server(
-    // C0
-	group: greet/s1
-)
-// C1
-service A-api {
-  // C2
-  @server( // C3
-    handler: GreetHandler
-  )
-  get /greet/from/:name(Request) returns (Response)   // hello
-	
-  // C4
-  @handler NoResponseHandler  // C5
-  get /greet/get(Request)
-}
-`
-
-const testMultiServiceTemplate = `
-info(
-    title: doc title
-    desc: doc description first part
-    version: 1.0
-)
-
-type Request struct {
-  Name string ` + "`" + `path:"name,options=you|me"` + "`" + ` 
-}
-
-type Response struct {
-  Message string ` + "`" + `json:"message"` + "`" + `
-}
-
-service A-api {
-  @server(
-    handler: GreetHandler
-  )
-  get /greet/from/:name(Request) returns (Response)
-}
-
-service A-api {
-  @server(
-    handler: NoResponseHandler
-  )
-  get /greet/get(Request)
-}
-`
-
-const apiNoInfo = `
-type Request struct {
-  Name string ` + "`" + `path:"name,options=you|me"` + "`" + `
-}
-
-type Response struct {
-  Message string ` + "`" + `json:"message"` + "`" + `
-}
-
-service A-api {
-  @server(
-    handler: GreetHandler
-  )
-  get /greet/from/:name(Request) returns (Response)
-}
-`
-
-const invalidApiFile = `
-type Request struct {
-  Name string ` + "`" + `path:"name,options=you|me"` + "`" + `
-}
-
-type Response struct {
-  Message string ` + "`" + `json:"message"` + "`" + `
-}
-
-service A-api
-  @server(
-    handler: GreetHandler
-  )
-  get /greet/from/:name(Request) returns (Response)
-}
-`
-
-const anonymousAnnotation = `
-type Request struct {
-  Name string ` + "`" + `path:"name,options=you|me"` + "`" + `
-}
-
-type Response struct {
-  Message string ` + "`" + `json:"message"` + "`" + `
-}
-
-service A-api {
-  @handler GreetHandler
-  get /greet/from/:name(Request) returns (Response)
-}
-`
-
-const apiHasMiddleware = `
-type Request struct {
-  Name string ` + "`" + `path:"name,options=you|me"` + "`" + `
-}
-
-type Response struct {
-  Message string ` + "`" + `json:"message"` + "`" + `
-}
-
-@server(
-	middleware: TokenValidate
-)
-service A-api {
-  @handler GreetHandler
-  get /greet/from/:name(Request) returns (Response)
-}
-`
-
-const apiJwt = `
-type Request struct {
-  Name string ` + "`" + `path:"name,options=you|me"` + "`" + `
-}
-
-type Response struct {
-  Message string ` + "`" + `json:"message"` + "`" + `
-}
-
-@server(
-	jwt: Auth
-	signature: true
-)
-service A-api {
-  @handler GreetHandler
-  get /greet/from/:name(Request) returns (Response)
-}
-`
-
-const apiJwtWithMiddleware = `
-type Request struct {
-  Name string ` + "`" + `path:"name,options=you|me"` + "`" + `
-}
-
-type Response struct {
-  Message string ` + "`" + `json:"message"` + "`" + `
-}
-
-@server(
-	jwt: Auth
-	jwtTransition: Trans
-	middleware: TokenValidate
-)
-service A-api {
-  @handler GreetHandler
-  get /greet/from/:name(Request) returns (Response)
-}
-`
-
-const apiHasNoRequest = `
-service A-api {
-  @handler GreetHandler
-  post /greet/ping ()
-}
-`
-
-const apiRouteTest = `
-type Request struct {
-  Name string ` + "`" + `path:"name,options=you|me"` + "`" + `
-}
-type Response struct {
-  Message string ` + "`" + `json:"message"` + "`" + `
-}
-service A-api {
-  @handler NormalHandler
-  get /greet/from/:name(Request) returns (Response)
-  @handler NoResponseHandler
-  get /greet/from/:sex(Request)
-  @handler NoRequestHandler
-  get /greet/from/request returns (Response)
-  @handler NoRequestNoResponseHandler
-  get /greet/from
-}
-`
-
-const hasCommentApiTest = `
-type Inline struct {
-
-}
-
-type Request struct {
-  Inline 
-  Name string ` + "`" + `path:"name,options=you|me"` + "`" + ` // name in path
-}
-
-type Response struct {
-  Message string ` + "`" + `json:"msg"` + "`" + ` // message
-}
-
-service A-api {
-  @doc ("helloworld")
-  @server(
-    handler: GreetHandler
-  )
-  get /greet/from/:name(Request) returns (Response)
-}
-`
-
-const hasInlineNoExistTest = `
-
-type Request struct {
-  Inline 
-  Name string ` + "`" + `path:"name,options=you|me"` + "`" + `
-}
-
-type Response struct {
-  Message string ` + "`" + `json:"message"` + "`" + ` // message
-}
-
-service A-api {
-  @doc ("helloworld")
-  @server(
-    handler: GreetHandler
-  )
-  get /greet/from/:name(Request) returns (Response)
-}
-`
-
-const importApi = `
-type ImportData struct {
-  Name string ` + "`" + `path:"name,options=you|me"` + "`" + `
-}
-
-`
-
-const hasImportApi = `
-import "importApi.api"
-
-type Request struct {
-  Name string ` + "`" + `path:"name,options=you|me"` + "`" + `
-}
-
-type Response struct {
-  Message string ` + "`" + `json:"message"` + "`" + ` // message
-}
-
-service A-api {
-  @server(
-    handler: GreetHandler
-  )
-  get /greet/from/:name(Request) returns (Response)
-}
-`
-
-const noStructTagApi = `
-type Request {
-  Name string ` + "`" + `path:"name,options=you|me"` + "`" + `
-}
-
-type XXX {}
-
-type (
-	Response {
-  		Message string ` + "`" + `json:"message"` + "`" + `
-	}
-
-	A {}
-
-	B struct {}
-)
-
-service A-api {
-  @handler GreetHandler
-  get /greet/from/:name(Request) returns (Response)
-}
-`
-
-const nestTypeApi = `
-type Request {
-  Name string ` + "`" + `path:"name,options=you|me"` + "`" + `
-  XXX struct {
-  }
-}
-
-service A-api {
-  @handler GreetHandler
-  get /greet/from/:name(Request)
-}
-`
-
 func TestParser(t *testing.T) {
 	filename := "greet.api"
 	err := ioutil.WriteFile(filename, []byte(testApiTemplate), os.ModePerm)
@@ -537,8 +269,15 @@ func validate(t *testing.T, api string) {
 }
 
 func validateWithCamel(t *testing.T, api, camel string) {
-	dir := t.TempDir()
-	err := DoGenProject(api, dir, camel)
+	dir := "workspace"
+	defer func() {
+		os.RemoveAll(dir)
+	}()
+	err := pathx.MkdirIfNotExist(dir)
+	assert.Nil(t, err)
+	err = initMod(dir)
+	assert.Nil(t, err)
+	err = DoGenProject(api, dir, camel)
 	assert.Nil(t, err)
 	filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
 		if strings.HasSuffix(path, ".go") {
@@ -550,6 +289,11 @@ func validateWithCamel(t *testing.T, api, camel string) {
 	})
 }
 
+func initMod(mod string) error {
+	_, err := execx.Run("go mod init "+mod, mod)
+	return err
+}
+
 func validateCode(code string) error {
 	_, err := goformat.Source([]byte(code))
 	return err

tools/goctl/api/gogen/genetc.go

@@ -1,6 +1,7 @@
 package gogen
 
 import (
+	_ "embed"
 	"fmt"
 	"strconv"
 
@@ -12,12 +13,11 @@ import (
 const (
 	defaultPort = 8888
 	etcDir      = "etc"
-	etcTemplate = `Name: {{.serviceName}}
-Host: {{.host}}
-Port: {{.port}}
-`
 )
 
+//go:embed etc.tpl
+var etcTemplate string
+
 func genEtc(dir string, cfg *config.Config, api *spec.ApiSpec) error {
 	filename, err := format.FileNamingFormat(cfg.NamingFormat, api.Service.Name)
 	if err != nil {

tools/goctl/api/gogen/genhandlers.go

@@ -1,6 +1,7 @@
 package gogen
 
 import (
+	_ "embed"
 	"fmt"
 	"path"
 	"strings"
@@ -14,36 +15,10 @@ import (
 	"github.com/zeromicro/go-zero/tools/goctl/vars"
 )
 
-const (
-	defaultLogicPackage = "logic"
-	handlerTemplate     = `package {{.PkgName}}
+const defaultLogicPackage = "logic"
 
-import (
-	"net/http"
-
-	{{if .After1_1_10}}"github.com/zeromicro/go-zero/rest/httpx"{{end}}
-	{{.ImportPackages}}
-)
-
-func {{.HandlerName}}(svcCtx *svc.ServiceContext) http.HandlerFunc {
-	return func(w http.ResponseWriter, r *http.Request) {
-		{{if .HasRequest}}var req types.{{.RequestType}}
-		if err := httpx.Parse(r, &req); err != nil {
-			httpx.Error(w, err)
-			return
-		}
-
-		{{end}}l := {{.LogicName}}.New{{.LogicType}}(r.Context(), svcCtx)
-		{{if .HasResp}}resp, {{end}}err := l.{{.Call}}({{if .HasRequest}}&req{{end}})
-		if err != nil {
-			httpx.Error(w, err)
-		} else {
-			{{if .HasResp}}httpx.OkJson(w, resp){{else}}httpx.Ok(w){{end}}
-		}
-	}
-}
-`
-)
+//go:embed handler.tpl
+var handlerTemplate string
 
 type handlerInfo struct {
 	PkgName        string
@@ -72,14 +47,10 @@ func genHandler(dir, rootPkg string, cfg *config.Config, group spec.Group, route
 		return err
 	}
 
-	goctlVersion := version.GetGoctlVersion()
-	// todo(anqiansong): This will be removed after a certain number of production versions of goctl (probably 5)
-	after1_1_10 := version.IsVersionGreaterThan(goctlVersion, "1.1.10")
 	return doGenToFile(dir, handler, cfg, group, route, handlerInfo{
 		PkgName:        pkgName,
 		ImportPackages: genHandlerImports(group, route, parentPkg),
 		HandlerName:    handler,
-		After1_1_10:    after1_1_10,
 		RequestType:    util.Title(route.RequestTypeName()),
 		LogicName:      logicName,
 		LogicType:      strings.Title(getLogicName(route)),

tools/goctl/api/gogen/genlogic.go

@@ -1,6 +1,7 @@
 package gogen
 
 import (
+	_ "embed"
 	"fmt"
 	"path"
 	"strconv"
@@ -14,32 +15,8 @@ import (
 	"github.com/zeromicro/go-zero/tools/goctl/vars"
 )
 
-const logicTemplate = `package {{.pkgName}}
-
-import (
-	{{.imports}}
-)
-
-type {{.logic}} struct {
-	logx.Logger
-	ctx    context.Context
-	svcCtx *svc.ServiceContext
-}
-
-func New{{.logic}}(ctx context.Context, svcCtx *svc.ServiceContext) *{{.logic}} {
-	return &{{.logic}}{
-		Logger: logx.WithContext(ctx),
-		ctx:    ctx,
-		svcCtx: svcCtx,
-	}
-}
-
-func (l *{{.logic}}) {{.function}}({{.request}}) {{.responseType}} {
-	// todo: add your logic here and delete this line
-
-	{{.returnString}}
-}
-`
+//go:embed logic.tpl
+var logicTemplate string
 
 func genLogic(dir, rootPkg string, cfg *config.Config, api *spec.ApiSpec) error {
 	for _, g := range api.Service.Groups {

tools/goctl/api/gogen/genmain.go

@@ -1,6 +1,7 @@
 package gogen
 
 import (
+	_ "embed"
 	"fmt"
 	"strings"
 
@@ -11,33 +12,8 @@ import (
 	"github.com/zeromicro/go-zero/tools/goctl/vars"
 )
 
-const mainTemplate = `package main
-
-import (
-	"flag"
-	"fmt"
-
-	{{.importPackages}}
-)
-
-var configFile = flag.String("f", "etc/{{.serviceName}}.yaml", "the config file")
-
-func main() {
-	flag.Parse()
-
-	var c config.Config
-	conf.MustLoad(*configFile, &c)
-
-	ctx := svc.NewServiceContext(c)
-	server := rest.MustNewServer(c.RestConf)
-	defer server.Stop()
-
-	handler.RegisterHandlers(server, ctx)
-
-	fmt.Printf("Starting server at %s:%d...\n", c.Host, c.Port)
-	server.Start()
-}
-`
+//go:embed main.tpl
+var mainTemplate string
 
 func genMain(dir, rootPkg string, cfg *config.Config, api *spec.ApiSpec) error {
 	name := strings.ToLower(api.Service.Name)

tools/goctl/api/gogen/genmiddleware.go

@@ -1,6 +1,7 @@
 package gogen
 
 import (
+	_ "embed"
 	"strings"
 
 	"github.com/zeromicro/go-zero/tools/goctl/api/spec"
@@ -8,27 +9,8 @@ import (
 	"github.com/zeromicro/go-zero/tools/goctl/util/format"
 )
 
-var middlewareImplementCode = `
-package middleware
-
-import "net/http"
-
-type {{.name}} struct {
-}
-
-func New{{.name}}() *{{.name}} {	
-	return &{{.name}}{}
-}
-
-func (m *{{.name}})Handle(next http.HandlerFunc) http.HandlerFunc {
-	return func(w http.ResponseWriter, r *http.Request) {
-		// TODO generate middleware implement function, delete after code implementation
-
-		// Passthrough to next handler if need 
-		next(w, r)
-	}	
-}
-`
+//go:embed middleware.tpl
+var middlewareImplementCode string
 
 func genMiddleware(dir string, cfg *config.Config, api *spec.ApiSpec) error {
 	middlewares := getMiddleware(api)

tools/goctl/api/gogen/genroutes.go

@@ -7,6 +7,7 @@ import (
 	"sort"
 	"strings"
 	"text/template"
+	"time"
 
 	"github.com/zeromicro/go-zero/core/collection"
 	"github.com/zeromicro/go-zero/tools/goctl/api/spec"
@@ -23,7 +24,8 @@ const (
 package handler
 
 import (
-	"net/http"
+	"net/http"{{if .hasTimeout}}
+	"time"{{end}}
 
 	{{.importPackages}}
 )
@@ -34,9 +36,10 @@ func RegisterHandlers(server *rest.Server, serverCtx *svc.ServiceContext) {
 `
 	routesAdditionTemplate = `
 	server.AddRoutes(
-		{{.routes}} {{.jwt}}{{.signature}} {{.prefix}}
+		{{.routes}} {{.jwt}}{{.signature}} {{.prefix}} {{.timeout}}
 	)
 `
+	timeoutThreshold = time.Millisecond
 )
 
 var mapping = map[string]string{
@@ -57,6 +60,7 @@ type (
 		jwtEnabled       bool
 		signatureEnabled bool
 		authName         string
+		timeout          string
 		middlewares      []string
 		prefix           string
 		jwtTrans         string
@@ -80,6 +84,7 @@ func genRoutes(dir, rootPkg string, cfg *config.Config, api *spec.ApiSpec) error
 		return err
 	}
 
+	var hasTimeout bool
 	gt := template.Must(template.New("groupTemplate").Parse(templateText))
 	for _, g := range groups {
 		var gbuilder strings.Builder
@@ -110,6 +115,22 @@ func genRoutes(dir, rootPkg string, cfg *config.Config, api *spec.ApiSpec) error
 rest.WithPrefix("%s"),`, g.prefix)
 		}
 
+		var timeout string
+		if len(g.timeout) > 0 {
+			duration, err := time.ParseDuration(g.timeout)
+			if err != nil {
+				return err
+			}
+
+			// why we check this, maybe some users set value 1, it's 1ns, not 1s.
+			if duration < timeoutThreshold {
+				return fmt.Errorf("timeout should not less than 1ms, now %v", duration)
+			}
+
+			timeout = fmt.Sprintf("rest.WithTimeout(%d * time.Millisecond),", duration/time.Millisecond)
+			hasTimeout = true
+		}
+
 		var routes string
 		if len(g.middlewares) > 0 {
 			gbuilder.WriteString("\n}...,")
@@ -130,6 +151,7 @@ rest.WithPrefix("%s"),`, g.prefix)
 			"jwt":       jwt,
 			"signature": signature,
 			"prefix":    prefix,
+			"timeout":   timeout,
 		}); err != nil {
 			return err
 		}
@@ -139,8 +161,8 @@ rest.WithPrefix("%s"),`, g.prefix)
 	if err != nil {
 		return err
 	}
-	routeFilename = routeFilename + ".go"
 
+	routeFilename = routeFilename + ".go"
 	filename := path.Join(dir, handlerDir, routeFilename)
 	os.Remove(filename)
 
@@ -152,7 +174,8 @@ rest.WithPrefix("%s"),`, g.prefix)
 		category:        category,
 		templateFile:    routesTemplateFile,
 		builtinTemplate: routesTemplate,
-		data: map[string]string{
+		data: map[string]interface{}{
+			"hasTimeout":      hasTimeout,
 			"importPackages":  genRouteImports(rootPkg, api),
 			"routesAdditions": strings.TrimSpace(builder.String()),
 		},
@@ -171,7 +194,8 @@ func genRouteImports(parentPkg string, api *spec.ApiSpec) string {
 					continue
 				}
 			}
-			importSet.AddStr(fmt.Sprintf("%s \"%s\"", toPrefix(folder), pathx.JoinPackages(parentPkg, handlerDir, folder)))
+			importSet.AddStr(fmt.Sprintf("%s \"%s\"", toPrefix(folder),
+				pathx.JoinPackages(parentPkg, handlerDir, folder)))
 		}
 	}
 	imports := importSet.KeysStr()
@@ -205,6 +229,8 @@ func getRoutes(api *spec.ApiSpec) ([]group, error) {
 			})
 		}
 
+		groupedRoutes.timeout = g.GetAnnotation("timeout")
+
 		jwt := g.GetAnnotation("jwt")
 		if len(jwt) > 0 {
 			groupedRoutes.authName = jwt

tools/goctl/api/gogen/gensvc.go

@@ -1,6 +1,7 @@
 package gogen
 
 import (
+	_ "embed"
 	"fmt"
 	"strings"
 
@@ -11,27 +12,10 @@ import (
 	"github.com/zeromicro/go-zero/tools/goctl/vars"
 )
 
-const (
-	contextFilename = "service_context"
-	contextTemplate = `package svc
+const contextFilename = "service_context"
 
-import (
-	{{.configImport}}
-)
-
-type ServiceContext struct {
-	Config {{.config}}
-	{{.middleware}}
-}
-
-func NewServiceContext(c {{.config}}) *ServiceContext {
-	return &ServiceContext{
-		Config: c, 
-		{{.middlewareAssignment}}
-	}
-}
-`
-)
+//go:embed svc.tpl
+var contextTemplate string
 
 func genServiceContext(dir, rootPkg string, cfg *config.Config, api *spec.ApiSpec) error {
 	filename, err := format.FileNamingFormat(cfg.NamingFormat, contextFilename)

tools/goctl/api/gogen/gentypes.go

@@ -1,6 +1,7 @@
 package gogen
 
 import (
+	_ "embed"
 	"fmt"
 	"io"
 	"os"
@@ -14,16 +15,10 @@ import (
 	"github.com/zeromicro/go-zero/tools/goctl/util/format"
 )
 
-const (
-	typesFile     = "types"
-	typesTemplate = `// Code generated by goctl. DO NOT EDIT.
-package types{{if .containsTime}}
-import (
-	"time"
-){{end}}
-{{.types}}
-`
-)
+const typesFile = "types"
+
+//go:embed types.tpl
+var typesTemplate string
 
 // BuildTypes gen types to string
 func BuildTypes(types []spec.Type) (string, error) {