feat/custom_buckets

rest/handler/prometheushandler.go

@@ -3,6 +3,7 @@ package handler
 import (
 	"net/http"
 	"strconv"
+	"sync"
 
 	"github.com/zeromicro/go-zero/core/metric"
 	"github.com/zeromicro/go-zero/core/timex"
@@ -12,14 +13,10 @@ 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},
-	})
+	rpcServerReqDurBuckets = []float64{5, 10, 25, 50, 100, 250, 500, 750, 1000}
+	metricServerReqDurOnce sync.Once
+
+	metricServerReqDur metric.HistogramVec
 
 	metricServerReqCodeTotal = metric.NewCounterVec(&metric.CounterVecOpts{
 		Namespace: serverNamespace,
@@ -30,8 +27,16 @@ var (
 	})
 )
 
+// SetServerReqDurBuckets sets buckets for rest server requests duration.
+// It must be called before PrometheusHandler is used.
+func SetServerReqDurBuckets(buckets []float64) {
+	rpcServerReqDurBuckets = buckets
+}
+
 // PrometheusHandler returns a middleware that reports stats to prometheus.
 func PrometheusHandler(path, method string) func(http.Handler) http.Handler {
+	initMetricServerReqDur()
+
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			startTime := timex.Now()
@@ -46,3 +51,16 @@ func PrometheusHandler(path, method string) func(http.Handler) http.Handler {
 		})
 	}
 }
+
+func initMetricServerReqDur() {
+	metricServerReqDurOnce.Do(func() {
+		metricServerReqDur = metric.NewHistogramVec(&metric.HistogramVecOpts{
+			Namespace: serverNamespace,
+			Subsystem: "requests",
+			Name:      "duration_ms",
+			Help:      "http server requests duration(ms).",
+			Labels:    []string{"path", "method", "code"},
+			Buckets:   rpcServerReqDurBuckets,
+		})
+	})
+}

rest/handler/prometheushandler_test.go

@@ -36,3 +36,8 @@ func TestPromMetricHandler_Enabled(t *testing.T) {
 	handler.ServeHTTP(resp, req)
 	assert.Equal(t, http.StatusOK, resp.Code)
 }
+
+func TestSetRpcServerReqDurBuckets(t *testing.T) {
+	SetServerReqDurBuckets([]float64{0.1})
+	assert.Equal(t, []float64{0.1}, rpcServerReqDurBuckets)
+}

zrpc/client.go

@@ -25,6 +25,9 @@ var (
 	WithTransportCredentials = internal.WithTransportCredentials
 	// WithUnaryClientInterceptor is an alias of internal.WithUnaryClientInterceptor.
 	WithUnaryClientInterceptor = internal.WithUnaryClientInterceptor
+	// SetRpcClientReqDurBuckets sets buckets for rpc client requests duration.
+	// It must be called before PrometheusInterceptor is used.
+	SetRpcClientReqDurBuckets = clientinterceptors.SetRpcClientReqDurBuckets
 )
 
 type (

zrpc/internal/clientinterceptors/prometheusinterceptor.go

@@ -3,6 +3,7 @@ package clientinterceptors
 import (
 	"context"
 	"strconv"
+	"sync"
 
 	"github.com/zeromicro/go-zero/core/metric"
 	"github.com/zeromicro/go-zero/core/timex"
@@ -13,14 +14,10 @@ import (
 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},
-	})
+	rpcClientReqDurBuckets = []float64{1, 2, 5, 10, 25, 50, 100, 250, 500, 1000, 2000, 5000}
+	metricClientReqDurOnce sync.Once
+
+	metricClientReqDur metric.HistogramVec
 
 	metricClientReqCodeTotal = metric.NewCounterVec(&metric.CounterVecOpts{
 		Namespace: clientNamespace,
@@ -31,12 +28,33 @@ var (
 	})
 )
 
+// SetRpcClientReqDurBuckets sets buckets for rpc client requests duration.
+// It must be called before PrometheusInterceptor is used.
+func SetRpcClientReqDurBuckets(buckets []float64) {
+	rpcClientReqDurBuckets = buckets
+}
+
 // 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 {
+	initMetricClientReqDur()
+
 	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 initMetricClientReqDur() {
+	metricClientReqDurOnce.Do(func() {
+		metricClientReqDur = metric.NewHistogramVec(&metric.HistogramVecOpts{
+			Namespace: clientNamespace,
+			Subsystem: "requests",
+			Name:      "duration_ms",
+			Help:      "rpc client requests duration(ms).",
+			Labels:    []string{"method"},
+			Buckets:   rpcClientReqDurBuckets,
+		})
+	})
+}

zrpc/internal/clientinterceptors/prometheusinterceptor_test.go

@@ -48,3 +48,9 @@ func TestPromMetricInterceptor(t *testing.T) {
 		})
 	}
 }
+
+func TestSetRpcClientReqDurBuckets(t *testing.T) {
+	buckets := []float64{1, 2, 3}
+	SetRpcClientReqDurBuckets(buckets)
+	assert.Equal(t, buckets, rpcClientReqDurBuckets)
+}

zrpc/internal/serverinterceptors/prometheusinterceptor.go

@@ -3,6 +3,7 @@ package serverinterceptors
 import (
 	"context"
 	"strconv"
+	"sync"
 
 	"github.com/zeromicro/go-zero/core/metric"
 	"github.com/zeromicro/go-zero/core/timex"
@@ -13,14 +14,10 @@ import (
 const serverNamespace = "rpc_server"
 
 var (
-	metricServerReqDur = metric.NewHistogramVec(&metric.HistogramVecOpts{
-		Namespace: serverNamespace,
-		Subsystem: "requests",
-		Name:      "duration_ms",
-		Help:      "rpc server requests duration(ms).",
-		Labels:    []string{"method"},
-		Buckets:   []float64{1, 2, 5, 10, 25, 50, 100, 250, 500, 1000, 2000, 5000},
-	})
+	rpcServerReqDurBuckets = []float64{1, 2, 5, 10, 25, 50, 100, 250, 500, 1000, 2000, 5000}
+	metricServerReqDurOnce sync.Once
+
+	metricServerReqDur metric.HistogramVec
 
 	metricServerReqCodeTotal = metric.NewCounterVec(&metric.CounterVecOpts{
 		Namespace: serverNamespace,
@@ -31,12 +28,33 @@ var (
 	})
 )
 
+// SetRpcServerReqDurBuckets sets buckets for rpc server requests duration.
+// It must be called before UnaryPrometheusInterceptor is used.
+func SetRpcServerReqDurBuckets(buckets []float64) {
+	rpcServerReqDurBuckets = buckets
+}
+
 // UnaryPrometheusInterceptor reports the statistics to the prometheus server.
 func UnaryPrometheusInterceptor(ctx context.Context, req any,
 	info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
+	initMetricServerReqDur()
+
 	startTime := timex.Now()
 	resp, err := handler(ctx, req)
 	metricServerReqDur.Observe(timex.Since(startTime).Milliseconds(), info.FullMethod)
 	metricServerReqCodeTotal.Inc(info.FullMethod, strconv.Itoa(int(status.Code(err))))
 	return resp, err
 }
+
+func initMetricServerReqDur() {
+	metricServerReqDurOnce.Do(func() {
+		metricServerReqDur = metric.NewHistogramVec(&metric.HistogramVecOpts{
+			Namespace: serverNamespace,
+			Subsystem: "requests",
+			Name:      "duration_ms",
+			Help:      "rpc server requests duration(ms).",
+			Labels:    []string{"method"},
+			Buckets:   rpcServerReqDurBuckets,
+		})
+	})
+}

zrpc/internal/serverinterceptors/prometheusinterceptor_test.go

@@ -30,3 +30,8 @@ func TestUnaryPromMetricInterceptor_Enabled(t *testing.T) {
 	})
 	assert.Nil(t, err)
 }
+
+func TestSetRpcServerReqDurBuckets(t *testing.T) {
+	SetRpcServerReqDurBuckets([]float64{0.1})
+	assert.Equal(t, []float64{0.1}, rpcServerReqDurBuckets)
+}

zrpc/server.go

@@ -13,6 +13,10 @@ import (
 	"google.golang.org/grpc"
 )
 
+// SetRpcServerReqDurBuckets sets buckets for rpc server requests duration.
+// It must be called before UnaryPrometheusInterceptor is used.
+var SetRpcServerReqDurBuckets = serverinterceptors.SetRpcServerReqDurBuckets
+
 // A RpcServer is a rpc server.
 type RpcServer struct {
 	server   internal.Server