diff --git a/core/stores/mongo/collection.go b/core/stores/mongo/collection.go index 1592dfa2..032289ac 100644 --- a/core/stores/mongo/collection.go +++ b/core/stores/mongo/collection.go @@ -11,7 +11,7 @@ import ( "github.com/tal-tech/go-zero/core/timex" ) -const slowThreshold = time.Millisecond * 500 +const defaultSlowThreshold = time.Millisecond * 500 // ErrNotFound is an alias of mgo.ErrNotFound. var ErrNotFound = mgo.ErrNotFound @@ -203,7 +203,7 @@ func (c *decoratedCollection) logDuration(method string, duration time.Duration, if e != nil { logx.Error(err) } else if err != nil { - if duration > slowThreshold { + if duration > slowThreshold.Load() { logx.WithDuration(duration).Slowf("[MONGO] mongo(%s) - slowcall - %s - fail(%s) - %s", c.name, method, err.Error(), string(content)) } else { @@ -211,7 +211,7 @@ func (c *decoratedCollection) logDuration(method string, duration time.Duration, c.name, method, err.Error(), string(content)) } } else { - if duration > slowThreshold { + if duration > slowThreshold.Load() { logx.WithDuration(duration).Slowf("[MONGO] mongo(%s) - slowcall - %s - ok - %s", c.name, method, string(content)) } else { diff --git a/core/stores/mongo/model.go b/core/stores/mongo/model.go index f42186d0..a13c1213 100644 --- a/core/stores/mongo/model.go +++ b/core/stores/mongo/model.go @@ -8,23 +8,14 @@ import ( "github.com/tal-tech/go-zero/core/breaker" ) -type ( - options struct { - timeout time.Duration - } - - // Option defines the method to customize a mongo model. - Option func(opts *options) - - // A Model is a mongo model. - Model struct { - session *concurrentSession - db *mgo.Database - collection string - brk breaker.Breaker - opts []Option - } -) +// A Model is a mongo model. +type Model struct { + session *concurrentSession + db *mgo.Database + collection string + brk breaker.Breaker + opts []Option +} // MustNewModel returns a Model, exits on errors. func MustNewModel(url, collection string, opts ...Option) *Model { diff --git a/core/stores/mongo/model_test.go b/core/stores/mongo/model_test.go new file mode 100644 index 00000000..f5cafe71 --- /dev/null +++ b/core/stores/mongo/model_test.go @@ -0,0 +1,14 @@ +package mongo + +import ( + "testing" + "time" + + "github.com/stretchr/testify/assert" +) + +func TestWithTimeout(t *testing.T) { + o := defaultOptions() + WithTimeout(time.Second)(o) + assert.Equal(t, time.Second, o.timeout) +} diff --git a/core/stores/mongo/options.go b/core/stores/mongo/options.go new file mode 100644 index 00000000..f90c77ea --- /dev/null +++ b/core/stores/mongo/options.go @@ -0,0 +1,28 @@ +package mongo + +import ( + "time" + + "github.com/tal-tech/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) +) + +func SetSlowThreshold(threshold time.Duration) { + slowThreshold.Set(threshold) +} + +func defaultOptions() *options { + return &options{ + timeout: defaultTimeout, + } +} diff --git a/core/stores/mongo/options_test.go b/core/stores/mongo/options_test.go new file mode 100644 index 00000000..372a642d --- /dev/null +++ b/core/stores/mongo/options_test.go @@ -0,0 +1,14 @@ +package mongo + +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()) +} diff --git a/core/stores/mongo/sessionmanager.go b/core/stores/mongo/sessionmanager.go index f262e321..b697b68c 100644 --- a/core/stores/mongo/sessionmanager.go +++ b/core/stores/mongo/sessionmanager.go @@ -57,9 +57,7 @@ func (cs *concurrentSession) putSession(session *mgo.Session) { } func (cs *concurrentSession) takeSession(opts ...Option) (*mgo.Session, error) { - o := &options{ - timeout: defaultTimeout, - } + o := defaultOptions() for _, opt := range opts { opt(o) }