refactor(zrpc): prefer static config for zrpc server statinterceptor (#3105)

zrpc/internal/config.go

@@ -1,5 +1,7 @@
 package internal
 
+import "time"
+
 type (
 	// ClientMiddlewaresConf defines whether to use client middlewares.
 	ClientMiddlewaresConf struct {
@@ -12,10 +14,12 @@ type (
 
 	// ServerMiddlewaresConf defines whether to use server middlewares.
 	ServerMiddlewaresConf struct {
-		Trace      bool `json:",default=true"`
-		Recover    bool `json:",default=true"`
-		Stat       bool `json:",default=true"`
-		Prometheus bool `json:",default=true"`
-		Breaker    bool `json:",default=true"`
+		Trace                    bool          `json:",default=true"`
+		Recover                  bool          `json:",default=true"`
+		Stat                     bool          `json:",default=true"`
+		StatSlowThreshold        time.Duration `json:",default=500ms"`
+		NotLoggingContentMethods []string      `json:",optional"`
+		Prometheus               bool          `json:",default=true"`
+		Breaker                  bool          `json:",default=true"`
 	}
 )

zrpc/internal/rpcserver.go

@@ -113,7 +113,11 @@ func (s *rpcServer) buildUnaryInterceptors() []grpc.UnaryServerInterceptor {
 		interceptors = append(interceptors, serverinterceptors.UnaryRecoverInterceptor)
 	}
 	if s.middlewares.Stat {
-		interceptors = append(interceptors, serverinterceptors.UnaryStatInterceptor(s.metrics))
+		interceptors = append(interceptors, serverinterceptors.UnaryStatInterceptor(s.metrics,
+			serverinterceptors.StatConf{
+				SlowThreshold:            s.middlewares.StatSlowThreshold,
+				NotLoggingContentMethods: s.middlewares.NotLoggingContentMethods,
+			}))
 	}
 	if s.middlewares.Prometheus {
 		interceptors = append(interceptors, serverinterceptors.UnaryPrometheusInterceptor)

zrpc/internal/serverinterceptors/statinterceptor.go

@@ -6,6 +6,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/zeromicro/go-zero/core/collection"
 	"github.com/zeromicro/go-zero/core/lang"
 	"github.com/zeromicro/go-zero/core/logx"
 	"github.com/zeromicro/go-zero/core/stat"
@@ -22,18 +23,29 @@ var (
 	slowThreshold            = syncx.ForAtomicDuration(defaultSlowThreshold)
 )
 
+// StatConf defines the static configuration for stat interceptor.
+type StatConf struct {
+	SlowThreshold            time.Duration
+	NotLoggingContentMethods []string
+}
+
 // DontLogContentForMethod disable logging content for given method.
+// Deprecated: use StatConf instead.
 func DontLogContentForMethod(method string) {
 	notLoggingContentMethods.Store(method, lang.Placeholder)
 }
 
 // SetSlowThreshold sets the slow threshold.
+// Deprecated: use StatConf instead.
 func SetSlowThreshold(threshold time.Duration) {
 	slowThreshold.Set(threshold)
 }
 
 // UnaryStatInterceptor returns a func that uses given metrics to report stats.
-func UnaryStatInterceptor(metrics *stat.Metrics) grpc.UnaryServerInterceptor {
+func UnaryStatInterceptor(metrics *stat.Metrics, conf StatConf) grpc.UnaryServerInterceptor {
+	staticNotLoggingContentMethods := collection.NewSet()
+	staticNotLoggingContentMethods.AddStr(conf.NotLoggingContentMethods...)
+
 	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo,
 		handler grpc.UnaryHandler) (resp any, err error) {
 		startTime := timex.Now()
@@ -42,14 +54,16 @@ func UnaryStatInterceptor(metrics *stat.Metrics) grpc.UnaryServerInterceptor {
 			metrics.Add(stat.Task{
 				Duration: duration,
 			})
-			logDuration(ctx, info.FullMethod, req, duration)
+			logDuration(ctx, info.FullMethod, req, duration,
+				staticNotLoggingContentMethods, conf.SlowThreshold)
 		}()
 
 		return handler(ctx, req)
 	}
 }
 
-func logDuration(ctx context.Context, method string, req any, duration time.Duration) {
+func logDuration(ctx context.Context, method string, req any, duration time.Duration,
+	staticNotLoggingContentMethods *collection.Set, staticSlowThreshold time.Duration) {
 	var addr string
 	client, ok := peer.FromContext(ctx)
 	if ok {
@@ -57,9 +71,8 @@ func logDuration(ctx context.Context, method string, req any, duration time.Dura
 	}
 
 	logger := logx.WithContext(ctx).WithDuration(duration)
-	_, ok = notLoggingContentMethods.Load(method)
-	if ok {
-		if duration > slowThreshold.Load() {
+	if !shouldLogContent(method, staticNotLoggingContentMethods) {
+		if isSlow(duration, staticSlowThreshold) {
 			logger.Slowf("[RPC] slowcall - %s - %s", addr, method)
 		}
 	} else {
@@ -73,3 +86,13 @@ func logDuration(ctx context.Context, method string, req any, duration time.Dura
 		}
 	}
 }
+
+func shouldLogContent(method string, staticNotLoggingContentMethods *collection.Set) bool {
+	_, ok := notLoggingContentMethods.Load(method)
+	return !ok && !staticNotLoggingContentMethods.Contains(method)
+}
+
+func isSlow(duration time.Duration, staticSlowThreshold time.Duration) bool {
+	return duration > slowThreshold.Load() ||
+		(staticSlowThreshold > 0 && duration > staticSlowThreshold)
+}

zrpc/internal/serverinterceptors/statinterceptor_test.go

@@ -3,12 +3,15 @@ package serverinterceptors
 import (
 	"context"
 	"net"
+	"sync"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/zeromicro/go-zero/core/collection"
 	"github.com/zeromicro/go-zero/core/lang"
 	"github.com/zeromicro/go-zero/core/stat"
+	"github.com/zeromicro/go-zero/core/syncx"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/peer"
 )
@@ -16,12 +19,16 @@ import (
 func TestSetSlowThreshold(t *testing.T) {
 	assert.Equal(t, defaultSlowThreshold, slowThreshold.Load())
 	SetSlowThreshold(time.Second)
+	// reset slowThreshold
+	t.Cleanup(func() {
+		slowThreshold = syncx.ForAtomicDuration(defaultSlowThreshold)
+	})
 	assert.Equal(t, time.Second, slowThreshold.Load())
 }
 
 func TestUnaryStatInterceptor(t *testing.T) {
 	metrics := stat.NewMetrics("mock")
-	interceptor := UnaryStatInterceptor(metrics)
+	interceptor := UnaryStatInterceptor(metrics, StatConf{})
 	_, err := interceptor(context.Background(), nil, &grpc.UnaryServerInfo{
 		FullMethod: "/",
 	}, func(ctx context.Context, req any) (any, error) {
@@ -78,7 +85,8 @@ func TestLogDuration(t *testing.T) {
 			t.Parallel()
 
 			assert.NotPanics(t, func() {
-				logDuration(test.ctx, "foo", test.req, test.duration)
+				logDuration(test.ctx, "foo", test.req, test.duration,
+					collection.NewSet(), 0)
 			})
 		})
 	}
@@ -127,14 +135,129 @@ func TestLogDurationWithoutContent(t *testing.T) {
 	}
 
 	DontLogContentForMethod("foo")
+	// reset notLoggingContentMethods
+	t.Cleanup(func() {
+		notLoggingContentMethods = sync.Map{}
+	})
 	for _, test := range tests {
 		test := test
 		t.Run(test.name, func(t *testing.T) {
 			t.Parallel()
 
 			assert.NotPanics(t, func() {
-				logDuration(test.ctx, "foo", test.req, test.duration)
+				logDuration(test.ctx, "foo", test.req, test.duration,
+					collection.NewSet(), 0)
+			})
+		})
+	}
+}
+
+func Test_shouldLogContent(t *testing.T) {
+	type args struct {
+		method                         string
+		staticNotLoggingContentMethods []string
+	}
+
+	tests := []struct {
+		name  string
+		args  args
+		want  bool
+		setup func()
+	}{
+		{
+			"empty",
+			args{
+				method: "foo",
+			},
+			true,
+			nil,
+		},
+		{
+			"static",
+			args{
+				method:                         "foo",
+				staticNotLoggingContentMethods: []string{"foo"},
+			},
+			false,
+			nil,
+		},
+		{
+			"dynamic",
+			args{
+				method: "foo",
+			},
+			false,
+			func() {
+				DontLogContentForMethod("foo")
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if tt.setup != nil {
+				tt.setup()
+			}
+			// reset notLoggingContentMethods
+			t.Cleanup(func() {
+				notLoggingContentMethods = sync.Map{}
+			})
+			set := collection.NewSet()
+			set.AddStr(tt.args.staticNotLoggingContentMethods...)
+			assert.Equalf(t, tt.want, shouldLogContent(tt.args.method, set), "shouldLogContent(%v, %v)", tt.args.method, tt.args.staticNotLoggingContentMethods)
+		})
+	}
+}
+
+func Test_isSlow(t *testing.T) {
+	type args struct {
+		duration            time.Duration
+		staticSlowThreshold time.Duration
+	}
+	tests := []struct {
+		name  string
+		args  args
+		want  bool
+		setup func()
+	}{
+		{
+			"default",
+			args{
+				duration: time.Millisecond * 501,
+			},
+			true,
+			nil,
+		},
+		{
+			"static",
+			args{
+				duration:            time.Millisecond * 200,
+				staticSlowThreshold: time.Millisecond * 100,
+			},
+			true,
+			nil,
+		},
+		{
+			"dynamic",
+			args{
+				duration: time.Millisecond * 200,
+			},
+			true,
+			func() {
+				SetSlowThreshold(time.Millisecond * 100)
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if tt.setup != nil {
+				tt.setup()
+			}
+			// reset slowThreshold
+			t.Cleanup(func() {
+				slowThreshold = syncx.ForAtomicDuration(defaultSlowThreshold)
 			})
+			assert.Equalf(t, tt.want, isSlow(tt.args.duration, tt.args.staticSlowThreshold), "isSlow(%v, %v)", tt.args.duration, tt.args.staticSlowThreshold)
 		})
 	}
 }

zrpc/server.go

@@ -100,11 +100,13 @@ func (rs *RpcServer) Stop() {
 }
 
 // DontLogContentForMethod disable logging content for given method.
+// Deprecated: use ServerMiddlewaresConf.NotLoggingContentMethods instead.
 func DontLogContentForMethod(method string) {
 	serverinterceptors.DontLogContentForMethod(method)
 }
 
 // SetServerSlowThreshold sets the slow threshold on server side.
+// Deprecated: use ServerMiddlewaresConf.SlowThreshold instead.
 func SetServerSlowThreshold(threshold time.Duration) {
 	serverinterceptors.SetSlowThreshold(threshold)
 }