Merge pull request #12 from AH-dark/feat/options

Tracer configuration using options

options.go

@@ -0,0 +1,105 @@
+/*
+ * Copyright 2021 CloudWeGo Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package prometheus
+
+import (
+	"net/http"
+
+	prom "github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/collectors"
+)
+
+var defaultBuckets = []float64{5000, 10000, 25000, 50000, 100000, 250000, 500000, 1000000}
+
+// Option opts for monitor prometheus
+type Option interface {
+	apply(cfg *config)
+}
+
+type option func(cfg *config)
+
+func (fn option) apply(cfg *config) {
+	fn(cfg)
+}
+
+type config struct {
+	buckets            []float64
+	enableGoCollector  bool
+	runtimeMetricRules []collectors.GoRuntimeMetricsRule
+	registry           *prom.Registry
+	serveMux           *http.ServeMux
+	disableServer      bool
+}
+
+func defaultConfig() *config {
+	return &config{
+		buckets:            defaultBuckets,
+		enableGoCollector:  false,
+		runtimeMetricRules: []collectors.GoRuntimeMetricsRule{},
+		registry:           prom.NewRegistry(),
+		serveMux:           http.DefaultServeMux,
+		disableServer:      false,
+	}
+}
+
+// WithEnableGoCollector enable go collector
+func WithEnableGoCollector(enable bool) Option {
+	return option(func(cfg *config) {
+		cfg.enableGoCollector = enable
+	})
+}
+
+// WithGoCollectorRule define your custom go collector rule
+func WithGoCollectorRule(rules ...collectors.GoRuntimeMetricsRule) Option {
+	return option(func(cfg *config) {
+		cfg.runtimeMetricRules = rules
+	})
+}
+
+// WithHistogramBuckets define your custom histogram buckets base on your biz
+func WithHistogramBuckets(buckets []float64) Option {
+	return option(func(cfg *config) {
+		if len(buckets) > 0 {
+			cfg.buckets = buckets
+		}
+	})
+}
+
+// WithRegistry define your custom registry
+func WithRegistry(registry *prom.Registry) Option {
+	return option(func(cfg *config) {
+		if registry != nil {
+			cfg.registry = registry
+		}
+	})
+}
+
+// WithServeMux define your custom serve mux
+func WithServeMux(serveMux *http.ServeMux) Option {
+	return option(func(cfg *config) {
+		if serveMux != nil {
+			cfg.serveMux = serveMux
+		}
+	})
+}
+
+// WithDisableServer disable prometheus server
+func WithDisableServer(disable bool) Option {
+	return option(func(cfg *config) {
+		cfg.disableServer = disable
+	})
+}

tracer.go

@@ -23,6 +23,7 @@ import (
 	"net/http"
 
 	prom "github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/collectors"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 
 	"github.com/cloudwego/kitex/pkg/rpcinfo"
@@ -100,14 +101,23 @@ func (c *clientTracer) Finish(ctx context.Context) {
 }
 
 // NewClientTracer provide tracer for client call, addr and path is the scrape_configs for prometheus server.
-func NewClientTracer(addr, path string) stats.Tracer {
-	registry := prom.NewRegistry()
-	http.Handle(path, promhttp.HandlerFor(registry, promhttp.HandlerOpts{ErrorHandling: promhttp.ContinueOnError}))
-	go func() {
-		if err := http.ListenAndServe(addr, nil); err != nil {
-			log.Fatal("Unable to start a promhttp server, err: " + err.Error())
-		}
-	}()
+func NewClientTracer(addr, path string, options ...Option) stats.Tracer {
+	cfg := defaultConfig()
+	for _, opt := range options {
+		opt.apply(cfg)
+	}
+
+	if !cfg.disableServer {
+		cfg.serveMux.Handle(path, promhttp.HandlerFor(cfg.registry, promhttp.HandlerOpts{
+			ErrorHandling: promhttp.ContinueOnError,
+			Registry:      cfg.registry,
+		}))
+		go func() {
+			if err := http.ListenAndServe(addr, cfg.serveMux); err != nil {
+				log.Fatal("Unable to start a promhttp server, err: " + err.Error())
+			}
+		}()
+	}
 
 	clientHandledCounter := prom.NewCounterVec(
 		prom.CounterOpts{
@@ -116,17 +126,21 @@ func NewClientTracer(addr, path string) stats.Tracer {
 		},
 		[]string{labelKeyCaller, labelKeyCallee, labelKeyMethod, labelKeyStatus, labelKeyRetry},
 	)
-	registry.MustRegister(clientHandledCounter)
+	cfg.registry.MustRegister(clientHandledCounter)
 
 	clientHandledHistogram := prom.NewHistogramVec(
 		prom.HistogramOpts{
 			Name:    "kitex_client_latency_us",
 			Help:    "Latency (microseconds) of the RPC until it is finished.",
-			Buckets: []float64{5000, 10000, 25000, 50000, 100000, 250000, 500000, 1000000},
+			Buckets: cfg.buckets,
 		},
 		[]string{labelKeyCaller, labelKeyCallee, labelKeyMethod, labelKeyStatus, labelKeyRetry},
 	)
-	registry.MustRegister(clientHandledHistogram)
+	cfg.registry.MustRegister(clientHandledHistogram)
+
+	if cfg.enableGoCollector {
+		cfg.registry.MustRegister(collectors.NewGoCollector(collectors.WithGoCollectorRuntimeMetrics(cfg.runtimeMetricRules...)))
+	}
 
 	return &clientTracer{
 		clientHandledCounter:   clientHandledCounter,
@@ -166,14 +180,23 @@ func (c *serverTracer) Finish(ctx context.Context) {
 }
 
 // NewServerTracer provides tracer for server access, addr and path is the scrape_configs for prometheus server.
-func NewServerTracer(addr, path string) stats.Tracer {
-	registry := prom.NewRegistry()
-	http.Handle(path, promhttp.HandlerFor(registry, promhttp.HandlerOpts{ErrorHandling: promhttp.ContinueOnError}))
-	go func() {
-		if err := http.ListenAndServe(addr, nil); err != nil {
-			log.Fatal("Unable to start a promhttp server, err: " + err.Error())
-		}
-	}()
+func NewServerTracer(addr, path string, options ...Option) stats.Tracer {
+	cfg := defaultConfig()
+	for _, opt := range options {
+		opt.apply(cfg)
+	}
+
+	if !cfg.disableServer {
+		cfg.serveMux.Handle(path, promhttp.HandlerFor(cfg.registry, promhttp.HandlerOpts{
+			ErrorHandling: promhttp.ContinueOnError,
+			Registry:      cfg.registry,
+		}))
+		go func() {
+			if err := http.ListenAndServe(addr, cfg.serveMux); err != nil {
+				log.Fatal("Unable to start a promhttp server, err: " + err.Error())
+			}
+		}()
+	}
 
 	serverHandledCounter := prom.NewCounterVec(
 		prom.CounterOpts{
@@ -182,17 +205,21 @@ func NewServerTracer(addr, path string) stats.Tracer {
 		},
 		[]string{labelKeyCaller, labelKeyCallee, labelKeyMethod, labelKeyStatus, labelKeyRetry},
 	)
-	registry.MustRegister(serverHandledCounter)
+	cfg.registry.MustRegister(serverHandledCounter)
 
 	serverHandledHistogram := prom.NewHistogramVec(
 		prom.HistogramOpts{
 			Name:    "kitex_server_latency_us",
 			Help:    "Latency (microseconds) of RPC that had been application-level handled by the server.",
-			Buckets: []float64{5000, 10000, 25000, 50000, 100000, 250000, 500000, 1000000},
+			Buckets: cfg.buckets,
 		},
 		[]string{labelKeyCaller, labelKeyCallee, labelKeyMethod, labelKeyStatus, labelKeyRetry},
 	)
-	registry.MustRegister(serverHandledHistogram)
+	cfg.registry.MustRegister(serverHandledHistogram)
+
+	if cfg.enableGoCollector {
+		cfg.registry.MustRegister(collectors.NewGoCollector(collectors.WithGoCollectorRuntimeMetrics(cfg.runtimeMetricRules...)))
+	}
 
 	return &serverTracer{
 		serverHandledCounter:   serverHandledCounter,