feat/custom_buckets

rest/engine.go

@@ -136,7 +136,7 @@ func (ng *engine) buildChainWithNativeMiddlewares(fr featuredRoutes, route Route
 		chn = chn.Append(ng.getLogHandler())
 	}
 	if ng.conf.Middlewares.Prometheus {
-		chn = chn.Append(handler.PrometheusHandler(route.Path, route.Method))
+		chn = chn.Append(handler.PrometheusHandler(route.Path, route.Method, fr.metricReqDurBuckets))
 	}
 	if ng.conf.Middlewares.MaxConns {
 		chn = chn.Append(handler.MaxConnsHandler(ng.conf.MaxConns))

rest/handler/prometheushandler.go

@@ -12,15 +12,7 @@ import (
 const serverNamespace = "http_server"
 
 var (
-	metricServerReqDur = metric.NewHistogramVec(&metric.HistogramVecOpts{
-		Namespace: serverNamespace,
-		Subsystem: "requests",
-		Name:      "duration_ms",
-		Help:      "http server requests duration(ms).",
-		Labels:    []string{"path", "method", "code"},
-		Buckets:   []float64{5, 10, 25, 50, 100, 250, 500, 750, 1000},
-	})
-
+	defaultDurationBuckets   = []float64{1, 2, 5, 10, 25, 50, 100, 250, 500, 1000, 2000, 5000}
 	metricServerReqCodeTotal = metric.NewCounterVec(&metric.CounterVecOpts{
 		Namespace: serverNamespace,
 		Subsystem: "requests",
@@ -31,18 +23,35 @@ var (
 )
 
 // PrometheusHandler returns a middleware that reports stats to prometheus.
-func PrometheusHandler(path, method string) func(http.Handler) http.Handler {
+func PrometheusHandler(path, method string, buckets []float64) func(http.Handler) http.Handler {
+	if len(buckets) == 0 {
+		buckets = defaultDurationBuckets
+	}
+
+	metricDurationHistogram := initMetricServerReqDur(buckets)
+
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			startTime := timex.Now()
 			cw := response.NewWithCodeResponseWriter(w)
 			defer func() {
 				code := strconv.Itoa(cw.Code)
-				metricServerReqDur.Observe(timex.Since(startTime).Milliseconds(), path, method, code)
+				metricDurationHistogram.Observe(timex.Since(startTime).Milliseconds(), path, method, code)
 				metricServerReqCodeTotal.Inc(path, method, code)
 			}()
 
 			next.ServeHTTP(cw, r)
 		})
 	}
 }
+
+func initMetricServerReqDur(buckets []float64) metric.HistogramVec {
+	return metric.NewHistogramVec(&metric.HistogramVecOpts{
+		Namespace: serverNamespace,
+		Subsystem: "requests",
+		Name:      "duration_ms",
+		Help:      "http server requests duration(ms).",
+		Labels:    []string{"path", "method", "code"},
+		Buckets:   buckets,
+	})
+}

rest/handler/prometheushandler_test.go

@@ -6,11 +6,12 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+
 	"github.com/zeromicro/go-zero/core/prometheus"
 )
 
 func TestPromMetricHandler_Disabled(t *testing.T) {
-	promMetricHandler := PrometheusHandler("/user/login", http.MethodGet)
+	promMetricHandler := PrometheusHandler("/user/login", http.MethodGet, nil)
 	handler := promMetricHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		w.WriteHeader(http.StatusOK)
 	}))
@@ -26,7 +27,7 @@ func TestPromMetricHandler_Enabled(t *testing.T) {
 		Host: "localhost",
 		Path: "/",
 	})
-	promMetricHandler := PrometheusHandler("/user/login", http.MethodGet)
+	promMetricHandler := PrometheusHandler("/user/login", http.MethodGet, nil)
 	handler := promMetricHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		w.WriteHeader(http.StatusOK)
 	}))

rest/server.go

@@ -219,6 +219,13 @@ func WithMaxBytes(maxBytes int64) RouteOption {
 	}
 }
 
+// WithMetricsReqDurBuckets returns a RouteOption to set metrics request duration buckets.
+func WithMetricsReqDurBuckets(buckets []float64) RouteOption {
+	return func(r *featuredRoutes) {
+		r.metricReqDurBuckets = buckets
+	}
+}
+
 // 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

@@ -15,6 +15,7 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/assert"
+
 	"github.com/zeromicro/go-zero/core/conf"
 	"github.com/zeromicro/go-zero/core/logx/logtest"
 	"github.com/zeromicro/go-zero/rest/chain"
@@ -134,6 +135,16 @@ func TestWithMaxBytes(t *testing.T) {
 	assert.Equal(t, int64(maxBytes), fr.maxBytes)
 }
 
+func TestWithMetricReqDurBuckets(t *testing.T) {
+	const (
+		bucket1 = 0.1
+		bucket2 = 0.5
+	)
+	var fr featuredRoutes
+	WithMetricsReqDurBuckets([]float64{bucket1, bucket2})(&fr)
+	assert.Equal(t, []float64{bucket1, bucket2}, fr.metricReqDurBuckets)
+}
+
 func TestWithMiddleware(t *testing.T) {
 	m := make(map[string]string)
 	rt := router.NewRouter()

rest/types.go

@@ -31,11 +31,12 @@ type (
 	}
 
 	featuredRoutes struct {
-		timeout   time.Duration
-		priority  bool
-		jwt       jwtSetting
-		signature signatureSetting
-		routes    []Route
-		maxBytes  int64
+		timeout             time.Duration
+		priority            bool
+		jwt                 jwtSetting
+		signature           signatureSetting
+		routes              []Route
+		maxBytes            int64
+		metricReqDurBuckets []float64
 	}
 )

zrpc/client.go

@@ -3,13 +3,14 @@ package zrpc
 import (
 	"time"
 
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/keepalive"
+
 	"github.com/zeromicro/go-zero/core/conf"
 	"github.com/zeromicro/go-zero/core/logx"
 	"github.com/zeromicro/go-zero/zrpc/internal"
 	"github.com/zeromicro/go-zero/zrpc/internal/auth"
 	"github.com/zeromicro/go-zero/zrpc/internal/clientinterceptors"
-	"google.golang.org/grpc"
-	"google.golang.org/grpc/keepalive"
 )
 
 var (
@@ -25,6 +26,8 @@ var (
 	WithTransportCredentials = internal.WithTransportCredentials
 	// WithUnaryClientInterceptor is an alias of internal.WithUnaryClientInterceptor.
 	WithUnaryClientInterceptor = internal.WithUnaryClientInterceptor
+	// WithClientMetricReqDurBuckets is an alias of internal.WithMetricReqDurBuckets.
+	WithClientMetricReqDurBuckets = internal.WithMetricReqDurBuckets
 )
 
 type (

zrpc/internal/client.go

@@ -7,12 +7,13 @@ import (
 	"strings"
 	"time"
 
-	"github.com/zeromicro/go-zero/zrpc/internal/balancer/p2c"
-	"github.com/zeromicro/go-zero/zrpc/internal/clientinterceptors"
-	"github.com/zeromicro/go-zero/zrpc/resolver"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/credentials/insecure"
+
+	"github.com/zeromicro/go-zero/zrpc/internal/balancer/p2c"
+	"github.com/zeromicro/go-zero/zrpc/internal/clientinterceptors"
+	"github.com/zeromicro/go-zero/zrpc/resolver"
 )
 
 const (
@@ -32,10 +33,11 @@ type (
 
 	// A ClientOptions is a client options.
 	ClientOptions struct {
-		NonBlock    bool
-		Timeout     time.Duration
-		Secure      bool
-		DialOptions []grpc.DialOption
+		NonBlock            bool
+		Timeout             time.Duration
+		Secure              bool
+		DialOptions         []grpc.DialOption
+		MetricReqDurBuckets []float64
 	}
 
 	// ClientOption defines the method to customize a ClientOptions.
@@ -84,7 +86,7 @@ func (c *client) buildDialOptions(opts ...ClientOption) []grpc.DialOption {
 	}
 
 	options = append(options,
-		grpc.WithChainUnaryInterceptor(c.buildUnaryInterceptors(cliOpts.Timeout)...),
+		grpc.WithChainUnaryInterceptor(c.buildUnaryInterceptors(cliOpts)...),
 		grpc.WithChainStreamInterceptor(c.buildStreamInterceptors()...),
 	)
 
@@ -101,7 +103,7 @@ func (c *client) buildStreamInterceptors() []grpc.StreamClientInterceptor {
 	return interceptors
 }
 
-func (c *client) buildUnaryInterceptors(timeout time.Duration) []grpc.UnaryClientInterceptor {
+func (c *client) buildUnaryInterceptors(cliOpts ClientOptions) []grpc.UnaryClientInterceptor {
 	var interceptors []grpc.UnaryClientInterceptor
 
 	if c.middlewares.Trace {
@@ -111,13 +113,13 @@ func (c *client) buildUnaryInterceptors(timeout time.Duration) []grpc.UnaryClien
 		interceptors = append(interceptors, clientinterceptors.DurationInterceptor)
 	}
 	if c.middlewares.Prometheus {
-		interceptors = append(interceptors, clientinterceptors.PrometheusInterceptor)
+		interceptors = append(interceptors, clientinterceptors.PrometheusInterceptor(cliOpts.MetricReqDurBuckets))
 	}
 	if c.middlewares.Breaker {
 		interceptors = append(interceptors, clientinterceptors.BreakerInterceptor)
 	}
 	if c.middlewares.Timeout {
-		interceptors = append(interceptors, clientinterceptors.TimeoutInterceptor(timeout))
+		interceptors = append(interceptors, clientinterceptors.TimeoutInterceptor(cliOpts.Timeout))
 	}
 
 	return interceptors
@@ -152,6 +154,13 @@ func WithDialOption(opt grpc.DialOption) ClientOption {
 	}
 }
 
+// WithMetricReqDurBuckets returns a func to customize a ClientOptions with given buckets.
+func WithMetricReqDurBuckets(buckets []float64) ClientOption {
+	return func(options *ClientOptions) {
+		options.MetricReqDurBuckets = buckets
+	}
+}
+
 // WithNonBlock sets the dialing to be nonblock.
 func WithNonBlock() ClientOption {
 	return func(options *ClientOptions) {

zrpc/internal/client_test.go

@@ -45,6 +45,13 @@ func TestWithStreamClientInterceptor(t *testing.T) {
 	assert.Equal(t, 1, len(options.DialOptions))
 }
 
+func TestWithMetricReqDurBuckets(t *testing.T) {
+	var options ClientOptions
+	opt := WithMetricReqDurBuckets([]float64{0.1, 0.2})
+	opt(&options)
+	assert.Equal(t, []float64{0.1, 0.2}, options.MetricReqDurBuckets)
+}
+
 func TestWithTransportCredentials(t *testing.T) {
 	var options ClientOptions
 	opt := WithTransportCredentials(nil)

zrpc/internal/clientinterceptors/prometheusinterceptor.go

@@ -4,24 +4,17 @@ import (
 	"context"
 	"strconv"
 
-	"github.com/zeromicro/go-zero/core/metric"
-	"github.com/zeromicro/go-zero/core/timex"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/status"
+
+	"github.com/zeromicro/go-zero/core/metric"
+	"github.com/zeromicro/go-zero/core/timex"
 )
 
 const clientNamespace = "rpc_client"
 
 var (
-	metricClientReqDur = metric.NewHistogramVec(&metric.HistogramVecOpts{
-		Namespace: clientNamespace,
-		Subsystem: "requests",
-		Name:      "duration_ms",
-		Help:      "rpc client requests duration(ms).",
-		Labels:    []string{"method"},
-		Buckets:   []float64{1, 2, 5, 10, 25, 50, 100, 250, 500, 1000, 2000, 5000},
-	})
-
+	defaultDurationBuckets   = []float64{1, 2, 5, 10, 25, 50, 100, 250, 500, 1000, 2000, 5000}
 	metricClientReqCodeTotal = metric.NewCounterVec(&metric.CounterVecOpts{
 		Namespace: clientNamespace,
 		Subsystem: "requests",
@@ -32,11 +25,30 @@ var (
 )
 
 // PrometheusInterceptor is an interceptor that reports to prometheus server.
-func PrometheusInterceptor(ctx context.Context, method string, req, reply any,
-	cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
-	startTime := timex.Now()
-	err := invoker(ctx, method, req, reply, cc, opts...)
-	metricClientReqDur.Observe(timex.Since(startTime).Milliseconds(), method)
-	metricClientReqCodeTotal.Inc(method, strconv.Itoa(int(status.Code(err))))
-	return err
+func PrometheusInterceptor(buckets []float64) grpc.UnaryClientInterceptor {
+	if len(buckets) == 0 {
+		buckets = defaultDurationBuckets
+	}
+
+	metricDurationHistogram := initMetricClientReqDur(buckets)
+
+	return func(ctx context.Context, method string, req, reply any,
+		cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
+		startTime := timex.Now()
+		err := invoker(ctx, method, req, reply, cc, opts...)
+		metricDurationHistogram.Observe(timex.Since(startTime).Milliseconds(), method)
+		metricClientReqCodeTotal.Inc(method, strconv.Itoa(int(status.Code(err))))
+		return err
+	}
+}
+
+func initMetricClientReqDur(buckets []float64) metric.HistogramVec {
+	return metric.NewHistogramVec(&metric.HistogramVecOpts{
+		Namespace: clientNamespace,
+		Subsystem: "requests",
+		Name:      "duration_ms",
+		Help:      "rpc client requests duration(ms).",
+		Labels:    []string{"method"},
+		Buckets:   buckets,
+	})
 }

zrpc/internal/clientinterceptors/prometheusinterceptor_test.go

@@ -6,8 +6,9 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/zeromicro/go-zero/core/prometheus"
 	"google.golang.org/grpc"
+
+	"github.com/zeromicro/go-zero/core/prometheus"
 )
 
 func TestPromMetricInterceptor(t *testing.T) {
@@ -39,7 +40,7 @@ func TestPromMetricInterceptor(t *testing.T) {
 				})
 			}
 			cc := new(grpc.ClientConn)
-			err := PrometheusInterceptor(context.Background(), "/foo", nil, nil, cc,
+			err := PrometheusInterceptor(nil)(context.Background(), "/foo", nil, nil, cc,
 				func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn,
 					opts ...grpc.CallOption) error {
 					return test.err