diff --git a/docs/_docs/dev-guide/tavern.md b/docs/_docs/dev-guide/tavern.md index 26f4eb21a..8ac998f85 100644 --- a/docs/_docs/dev-guide/tavern.md +++ b/docs/_docs/dev-guide/tavern.md @@ -14,6 +14,15 @@ If you would like to help contribute to Tavern, please take a look at our [open ## Configuration +### Metrics + +By default, Tavern does not export metrics. You may use the below environment configuration variables to enable [Prometheus](https://prometheus.io/docs/introduction/overview/) metric collection. These metrics become available at the "/metrics" endpoint configured. These metrics are hosted on a separate HTTP server such that it can be restricted to localhost (default). This is because the endpoint is unauthenticated, and would leak sensitive information if it was accessible. + +| Env Var | Description | Default | Required | +| ------- | ----------- | ------- | -------- | +| ENABLE_METRICS | Set to any value to enable the "/metrics" endpoint. | Disabled | No | +| HTTP_METRICS_LISTEN_ADDR | Listen address for the metrics HTTP server, it must be different than the value of `HTTP_LISTEN_ADDR`. | `127.0.0.1:8080` | No | + ### MySQL By default, Tavern operates an in-memory SQLite database. To persist data, a MySQL backend is supported. In order to configure Tavern to use MySQL, the `MYSQL_ADDR` environment variable must be set to the `host[:port]` of the database (e.g. `127.0.0.1`, `mydb.com`, or `mydb.com:3306`). You can reference the [mysql.Config](https://pkg.go.dev/github.com/go-sql-driver/mysql#Config) for additional information about Tavern's MySQL configuration. diff --git a/go.mod b/go.mod index 8235cef85..1b3a04926 100644 --- a/go.mod +++ b/go.mod @@ -17,7 +17,7 @@ require ( github.com/vektah/gqlparser/v2 v2.5.10 golang.org/x/crypto v0.14.0 golang.org/x/net v0.17.0 - golang.org/x/oauth2 v0.11.0 + golang.org/x/oauth2 v0.12.0 golang.org/x/sync v0.4.0 google.golang.org/grpc v1.59.0 google.golang.org/protobuf v1.32.0 @@ -31,6 +31,8 @@ require ( github.com/agext/levenshtein v1.2.3 // indirect github.com/agnivade/levenshtein v1.1.1 // indirect github.com/apparentlymart/go-textseg/v15 v15.0.0 // indirect + github.com/beorn7/perks v1.0.1 // indirect + github.com/cespare/xxhash/v2 v2.2.0 // indirect github.com/cpuguy83/go-md2man/v2 v2.0.2 // indirect github.com/davecgh/go-spew v1.1.1 // indirect github.com/go-openapi/inflect v0.19.0 // indirect @@ -39,9 +41,14 @@ require ( github.com/hashicorp/errwrap v1.1.0 // indirect github.com/hashicorp/golang-lru/v2 v2.0.3 // indirect github.com/hashicorp/hcl/v2 v2.18.1 // indirect + github.com/matttproud/golang_protobuf_extensions/v2 v2.0.0 // indirect github.com/mitchellh/go-wordwrap v1.0.1 // indirect github.com/mitchellh/mapstructure v1.5.0 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect + github.com/prometheus/client_golang v1.18.0 // indirect + github.com/prometheus/client_model v0.5.0 // indirect + github.com/prometheus/common v0.45.0 // indirect + github.com/prometheus/procfs v0.12.0 // indirect github.com/russross/blackfriday/v2 v2.1.0 // indirect github.com/sosodev/duration v1.2.0 // indirect github.com/vmihailenco/msgpack/v5 v5.4.0 // indirect @@ -49,7 +56,7 @@ require ( github.com/zclconf/go-cty v1.14.1 // indirect golang.org/x/exp v0.0.0-20231006140011-7918f672742d // indirect golang.org/x/mod v0.13.0 // indirect - golang.org/x/sys v0.13.0 // indirect + golang.org/x/sys v0.15.0 // indirect golang.org/x/text v0.13.0 // indirect golang.org/x/tools v0.14.0 // indirect google.golang.org/appengine v1.6.7 // indirect diff --git a/go.sum b/go.sum index 3987fea1c..cbbd0abfc 100644 --- a/go.sum +++ b/go.sum @@ -21,6 +21,10 @@ github.com/apparentlymart/go-textseg/v15 v15.0.0 h1:uYvfpb3DyLSCGWnctWKGj857c6ew github.com/apparentlymart/go-textseg/v15 v15.0.0/go.mod h1:K8XmNZdhEBkdlyDdvbmmsvpAG721bKi0joRfFdHIWJ4= github.com/arbovm/levenshtein v0.0.0-20160628152529-48b4e1c0c4d0 h1:jfIu9sQUG6Ig+0+Ap1h4unLjW6YQJpKZVmUzxsD4E/Q= github.com/arbovm/levenshtein v0.0.0-20160628152529-48b4e1c0c4d0/go.mod h1:t2tdKJDJF9BV14lnkjHmOQgcvEKgtqs5a1N3LNdJhGE= +github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= +github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw= +github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44= +github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/cpuguy83/go-md2man/v2 v2.0.0-20190314233015-f79a8a8ca69d/go.mod h1:maD7wRr/U5Z6m/iR4s+kqSMx2CaBsrgA7czyZG/E6dU= github.com/cpuguy83/go-md2man/v2 v2.0.2 h1:p1EgwI/C7NhT0JmVkwCD2ZBK8j4aeHQX2pMHHBfMQ6w= github.com/cpuguy83/go-md2man/v2 v2.0.2/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o= @@ -57,16 +61,27 @@ github.com/hashicorp/golang-lru/v2 v2.0.3/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyf github.com/hashicorp/hcl/v2 v2.18.1 h1:6nxnOJFku1EuSawSD81fuviYUV8DxFr3fp2dUi3ZYSo= github.com/hashicorp/hcl/v2 v2.18.1/go.mod h1:ThLC89FV4p9MPW804KVbe/cEXoQ8NZEh+JtMeeGErHE= github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI= +github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kylelemons/godebug v0.0.0-20170820004349-d65d576e9348 h1:MtvEpTB6LX3vkb4ax0b5D2DHbNAUsen0Gx5wZoq3lV4= github.com/mattn/go-sqlite3 v1.14.16 h1:yOQRA0RpS5PFz/oikGwBEqvAWhWg5ufRz4ETLjwpU1Y= github.com/mattn/go-sqlite3 v1.14.16/go.mod h1:2eHXhiwb8IkHr+BDWZGa96P6+rkvnG63S2DGjv9HUNg= +github.com/matttproud/golang_protobuf_extensions/v2 v2.0.0 h1:jWpvCLoY8Z/e3VKvlsiIGKtc+UG6U5vzxaoagmhXfyg= +github.com/matttproud/golang_protobuf_extensions/v2 v2.0.0/go.mod h1:QUyp042oQthUoa9bqDv0ER0wrtXnBruoNd7aNjkbP+k= github.com/mitchellh/go-wordwrap v1.0.1 h1:TLuKupo69TCn6TQSyGxwI1EblZZEsQ0vMlAFQflz0v0= github.com/mitchellh/go-wordwrap v1.0.1/go.mod h1:R62XHJLzvMFRBbcrT7m7WgmE1eOyTSsCt+hzestvNj0= github.com/mitchellh/mapstructure v1.5.0 h1:jeMsZIYE/09sWLaz43PL7Gy6RuMjD2eJVyuac5Z2hdY= github.com/mitchellh/mapstructure v1.5.0/go.mod h1:bFUtVrKA4DC2yAKiSyO/QUcy7e+RRV2QTWOzhPopBRo= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/prometheus/client_golang v1.18.0 h1:HzFfmkOzH5Q8L8G+kSJKUx5dtG87sewO+FoDDqP5Tbk= +github.com/prometheus/client_golang v1.18.0/go.mod h1:T+GXkCk5wSJyOqMIzVgvvjFDlkOQntgjkJWKrN5txjA= +github.com/prometheus/client_model v0.5.0 h1:VQw1hfvPvk3Uv6Qf29VrPF32JB6rtbgI6cYPYQjL0Qw= +github.com/prometheus/client_model v0.5.0/go.mod h1:dTiFglRmd66nLR9Pv9f0mZi7B7fk5Pm3gvsjB5tr+kI= +github.com/prometheus/common v0.45.0 h1:2BGz0eBc2hdMDLnO/8n0jeB3oPrt2D08CekT0lneoxM= +github.com/prometheus/common v0.45.0/go.mod h1:YJmSTw9BoKxJplESWWxlbyttQR4uaEcGyv9MZjVOJsY= +github.com/prometheus/procfs v0.12.0 h1:jluTpSng7V9hY0O2R9DzzJHYb2xULk9VTR1V1R/k6Bo= +github.com/prometheus/procfs v0.12.0/go.mod h1:pcuDEFsWDnvcgNzo4EEweacyhjeA9Zk3cnaOZAZEfOo= github.com/russross/blackfriday/v2 v2.0.1/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= github.com/russross/blackfriday/v2 v2.1.0 h1:JIOH55/0cWyOuilr9/qlrm0BSXldqnqwMsf35Ld67mk= github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= @@ -104,11 +119,15 @@ golang.org/x/net v0.17.0 h1:pVaXccu2ozPjCXewfr1S7xza/zcXTity9cCdXQYSjIM= golang.org/x/net v0.17.0/go.mod h1:NxSsAGuq816PNPmqtQdLE42eU2Fs7NoRIZrHJAlaCOE= golang.org/x/oauth2 v0.11.0 h1:vPL4xzxBM4niKCW6g9whtaWVXTJf1U5e4aZxxFx/gbU= golang.org/x/oauth2 v0.11.0/go.mod h1:LdF7O/8bLR/qWK9DrpXmbHLTouvRHK0SgJl0GmDBchk= +golang.org/x/oauth2 v0.12.0 h1:smVPGxink+n1ZI5pkQa8y6fZT0RW0MgCO5bFpepy4B4= +golang.org/x/oauth2 v0.12.0/go.mod h1:A74bZ3aGXgCY0qaIC9Ahg6Lglin4AMAco8cIv9baba4= golang.org/x/sync v0.4.0 h1:zxkM55ReGkDlKSM+Fu41A+zmbZuaPVbGMzvvdUPznYQ= golang.org/x/sync v0.4.0/go.mod h1:FU7BRWz2tNW+3quACPkgCx/L+uEAv1htQ0V83Z9Rj+Y= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.13.0 h1:Af8nKPmuFypiUBjVoU9V20FiaFXOcuZI21p0ycVYYGE= golang.org/x/sys v0.13.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.15.0 h1:h48lPFYpsTvQJZF4EKyI4aLHaev3CxivZmv7yZig9pc= +golang.org/x/sys v0.15.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk= golang.org/x/text v0.13.0 h1:ablQoSUd0tRdKxZewP80B+BaqeKJuVhuRxj/dkrun3k= diff --git a/tavern/app.go b/tavern/app.go index 4a680e86d..d7471af89 100644 --- a/tavern/app.go +++ b/tavern/app.go @@ -8,29 +8,29 @@ import ( "fmt" "log" "net/http" - pprof "net/http/pprof" + "net/http/pprof" "os" "strings" "entgo.io/contrib/entgql" - "google.golang.org/grpc" - "realm.pub/tavern/internal/c2" - "realm.pub/tavern/internal/c2/c2pb" - "realm.pub/tavern/internal/graphql" - "realm.pub/tavern/tomes" - gqlgraphql "github.com/99designs/gqlgen/graphql" "github.com/99designs/gqlgen/graphql/handler" "github.com/99designs/gqlgen/graphql/playground" + "github.com/prometheus/client_golang/prometheus/promhttp" "github.com/urfave/cli" "golang.org/x/net/http2" "golang.org/x/net/http2/h2c" + "google.golang.org/grpc" "realm.pub/tavern/internal/auth" + "realm.pub/tavern/internal/c2" + "realm.pub/tavern/internal/c2/c2pb" "realm.pub/tavern/internal/cdn" "realm.pub/tavern/internal/ent" "realm.pub/tavern/internal/ent/migrate" + "realm.pub/tavern/internal/graphql" tavernhttp "realm.pub/tavern/internal/http" "realm.pub/tavern/internal/www" + "realm.pub/tavern/tomes" ) func newApp(ctx context.Context, options ...func(*Config)) (app *cli.App) { @@ -50,7 +50,24 @@ func run(ctx context.Context, options ...func(*Config)) error { if err != nil { return err } - defer srv.client.Close() + defer srv.Close() + + // Start Metrics Server (if configured) + if srv.MetricsHTTP != nil { + if srv.HTTP.Addr == srv.MetricsHTTP.Addr { + return fmt.Errorf( + "tavern and metrics http must have different listen configurations (tavern=%q, metrics=%q)", + srv.HTTP.Addr, + srv.MetricsHTTP.Addr, + ) + } + go func() { + log.Printf("Metrics HTTP Server started on %s", srv.MetricsHTTP.Addr) + if err := srv.MetricsHTTP.ListenAndServe(); err != nil { + log.Printf("[WARN] stopped metrics http server: %v", err) + } + }() + } // Listen & Serve HTTP Traffic log.Printf("Starting HTTP server on %s", srv.HTTP.Addr) @@ -63,13 +80,15 @@ func run(ctx context.Context, options ...func(*Config)) error { // Server responsible for handling Tavern requests. type Server struct { - HTTP *http.Server - client *ent.Client + HTTP *http.Server + MetricsHTTP *http.Server + client *ent.Client } // Close should always be called to clean up a Tavern server. func (srv *Server) Close() error { srv.HTTP.Shutdown(context.Background()) + srv.MetricsHTTP.Shutdown(context.Background()) return srv.client.Close() } @@ -184,6 +203,12 @@ func NewServer(ctx context.Context, options ...func(*Config)) (*Server, error) { client: client, } + // Setup Metrics + if cfg.IsMetricsEnabled() { + log.Printf("[WARN] Metrics reporting is enabled, unauthenticated /metrics endpoint will be available at %q", EnvHTTPMetricsListenAddr.String()) + tSrv.MetricsHTTP = newMetricsServer() + } + // Shutdown for Test Run & Exit if cfg.IsTestRunAndExitEnabled() { go func() { @@ -227,7 +252,10 @@ func newGraphQLHandler(client *ent.Client) http.Handler { func newGRPCHandler(client *ent.Client) http.Handler { c2srv := c2.New(client) - grpcSrv := grpc.NewServer() + grpcSrv := grpc.NewServer( + grpc.UnaryInterceptor(grpcWithUnaryMetrics), + grpc.StreamInterceptor(grpcWithStreamMetrics), + ) c2pb.RegisterC2Server(grpcSrv, c2srv) return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.ProtoMajor != 2 { @@ -244,6 +272,15 @@ func newGRPCHandler(client *ent.Client) http.Handler { }) } +func newMetricsServer() *http.Server { + router := http.NewServeMux() + router.Handle("/metrics", promhttp.Handler()) + return &http.Server{ + Addr: EnvHTTPMetricsListenAddr.String(), + Handler: router, + } +} + func registerProfiler(router tavernhttp.RouteMap) { router.HandleFunc("/debug/pprof/", pprof.Index) router.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline) diff --git a/tavern/config.go b/tavern/config.go index 88f150d03..2c93488d5 100644 --- a/tavern/config.go +++ b/tavern/config.go @@ -22,7 +22,9 @@ var ( EnvEnableTestRunAndExit = EnvString{"ENABLE_TEST_RUN_AND_EXIT", ""} // EnvHTTPListenAddr sets the address (ip:port) for tavern's HTTP server to bind to. - EnvHTTPListenAddr = EnvString{"HTTP_LISTEN_ADDR", "0.0.0.0:80"} + // EnvHTTPMetricsAddr sets the address (ip:port) for the HTTP metrics server to bind to. + EnvHTTPListenAddr = EnvString{"HTTP_LISTEN_ADDR", "0.0.0.0:80"} + EnvHTTPMetricsListenAddr = EnvString{"HTTP_METRICS_LISTEN_ADDR", "127.0.0.1:8080"} // EnvOAuthClientID set to configure OAuth Client ID. // EnvOAuthClientSecret set to configure OAuth Client Secret. @@ -50,7 +52,9 @@ var ( EnvDBMaxConnLifetime = EnvInteger{"DB_MAX_CONN_LIFETIME", 3600} // EnvEnablePProf enables performance profiling and should not be enabled in production. - EnvEnablePProf = EnvString{"ENABLE_PPROF", ""} + // EnvEnableMetrics enables the /metrics endpoint and HTTP server. It is unauthenticated and should be used carefully. + EnvEnablePProf = EnvString{"ENABLE_PPROF", ""} + EnvEnableMetrics = EnvString{"ENABLE_METRICS", ""} ) // Config holds information that controls the behaviour of Tavern @@ -107,6 +111,11 @@ func (cfg *Config) Connect(options ...ent.Option) (*ent.Client, error) { return ent.NewClient(append(options, ent.Driver(drv))...), nil } +// IsMetricsEnabled returns true if the /metrics http endpoint has been enabled. +func (cfg *Config) IsMetricsEnabled() bool { + return EnvEnableMetrics.String() != "" +} + // IsPProfEnabled returns true if performance profiling has been enabled. func (cfg *Config) IsPProfEnabled() bool { return EnvEnablePProf.String() != "" diff --git a/tavern/internal/c2/api_claim_tasks.go b/tavern/internal/c2/api_claim_tasks.go index 0a1c147fa..f5d4baa81 100644 --- a/tavern/internal/c2/api_claim_tasks.go +++ b/tavern/internal/c2/api_claim_tasks.go @@ -6,6 +6,7 @@ import ( "fmt" "time" + "github.com/prometheus/client_golang/prometheus" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" "realm.pub/tavern/internal/c2/c2pb" @@ -14,6 +15,20 @@ import ( "realm.pub/tavern/internal/namegen" ) +var ( + metricHostCallbacksTotal = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Name: "tavern_host_callbacks_total", + Help: "The total number of ClaimTasks gRPC calls, provided with host labeling", + }, + []string{"host_identifier"}, + ) +) + +func init() { + prometheus.MustRegister(metricHostCallbacksTotal) +} + func (srv *Server) ClaimTasks(ctx context.Context, req *c2pb.ClaimTasksRequest) (*c2pb.ClaimTasksResponse, error) { now := time.Now() @@ -40,6 +55,11 @@ func (srv *Server) ClaimTasks(ctx context.Context, req *c2pb.ClaimTasksRequest) return nil, status.Errorf(codes.InvalidArgument, "must provide agent identifier") } + // Metrics + defer func() { + metricHostCallbacksTotal.WithLabelValues(req.Beacon.Identifier).Inc() + }() + // Upsert the host hostID, err := srv.graph.Host.Create(). SetIdentifier(req.Beacon.Host.Identifier). diff --git a/tavern/internal/http/metrics.go b/tavern/internal/http/metrics.go new file mode 100644 index 000000000..ec0fc1015 --- /dev/null +++ b/tavern/internal/http/metrics.go @@ -0,0 +1,27 @@ +package http + +import "github.com/prometheus/client_golang/prometheus" + +var ( + metricHTTPRequests = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Name: "tavern_http_requests_total", + Help: "Total number of requests received.", + }, + []string{"request_uri", "method"}, + ) + + metricHTTPLatency = prometheus.NewHistogramVec( + prometheus.HistogramOpts{ + Name: "tavern_http_request_duration_seconds", + Help: "Latency of requests.", + Buckets: prometheus.DefBuckets, + }, + []string{"request_uri", "method"}, + ) +) + +func init() { + // Register metrics with Prometheus + prometheus.MustRegister(metricHTTPRequests, metricHTTPLatency) +} diff --git a/tavern/internal/http/server.go b/tavern/internal/http/server.go index 11733ebf8..ba07fa8c4 100644 --- a/tavern/internal/http/server.go +++ b/tavern/internal/http/server.go @@ -2,6 +2,7 @@ package http import ( "net/http" + "time" ) // A Server for Tavern HTTP traffic. @@ -12,6 +13,8 @@ type Server struct { } func (srv *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) { + start := time.Now() + // Authenticate Request (if possible) ctx, err := srv.Authenticate(r) if err != nil { @@ -32,6 +35,15 @@ func (srv *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) { } r = r.WithContext(ctx) + // Metrics + defer func() { + // Increment total requests counter + metricHTTPRequests.WithLabelValues(r.RequestURI, r.Method).Inc() + + // Record the latency + metricHTTPLatency.WithLabelValues(r.RequestURI, r.Method).Observe(time.Since(start).Seconds()) + }() + // Log Request if srv.Logger != nil { srv.Log(r) diff --git a/tavern/main_test.go b/tavern/main_test.go index bc8e6d37d..c72d524f8 100644 --- a/tavern/main_test.go +++ b/tavern/main_test.go @@ -10,12 +10,16 @@ import ( func TestMainFunc(t *testing.T) { os.Setenv(EnvEnableTestRunAndExit.Key, "1") os.Setenv(EnvHTTPListenAddr.Key, "127.0.0.1:8080") + os.Setenv(EnvHTTPMetricsListenAddr.Key, "127.0.0.1:8081") os.Setenv(EnvEnablePProf.Key, "1") + os.Setenv(EnvEnableMetrics.Key, "1") defer func() { unsetList := []string{ EnvEnableTestRunAndExit.Key, EnvHTTPListenAddr.Key, + EnvHTTPMetricsListenAddr.Key, EnvEnablePProf.Key, + EnvEnableMetrics.Key, } for _, unset := range unsetList { if err := os.Unsetenv(unset); err != nil { diff --git a/tavern/metrics.go b/tavern/metrics.go new file mode 100644 index 000000000..90cf6d30d --- /dev/null +++ b/tavern/metrics.go @@ -0,0 +1,102 @@ +package main + +import ( + "context" + "time" + + "github.com/prometheus/client_golang/prometheus" + "google.golang.org/grpc" +) + +var ( + metricGRPCRequests = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Name: "tavern_grpc_requests_total", + Help: "Total number of requests received.", + }, + []string{"method"}, + ) + + metricGRPCLatency = prometheus.NewHistogramVec( + prometheus.HistogramOpts{ + Name: "tavern_grpc_request_duration_seconds", + Help: "Latency of requests.", + Buckets: prometheus.DefBuckets, + }, + []string{"method"}, + ) + + metricGRPCErrors = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Name: "tavern_grpc_request_errors", + Help: "Total number of errors.", + }, + []string{"method"}, + ) +) + +func init() { + // Register metrics with Prometheus + prometheus.MustRegister(metricGRPCRequests, metricGRPCLatency, metricGRPCErrors) +} + +func grpcWithUnaryMetrics( + ctx context.Context, + req interface{}, + info *grpc.UnaryServerInfo, + handler grpc.UnaryHandler, +) (interface{}, error) { + var ( + start = time.Now() + h any + err error + ) + + defer func() { + // Increment total requests counter + metricGRPCRequests.WithLabelValues(info.FullMethod).Inc() + + // Record the latency + metricGRPCLatency.WithLabelValues(info.FullMethod).Observe(time.Since(start).Seconds()) + + // Record if there was an error + if err != nil { + metricGRPCErrors.WithLabelValues(info.FullMethod).Inc() + } + }() + + // Call the handler + h, err = handler(ctx, req) + + return h, err +} + +func grpcWithStreamMetrics( + srv interface{}, + ss grpc.ServerStream, + info *grpc.StreamServerInfo, + handler grpc.StreamHandler, +) error { + var ( + start = time.Now() + err error + ) + + defer func() { + // Increment total requests counter + metricGRPCRequests.WithLabelValues(info.FullMethod).Inc() + + // Record the latency + metricGRPCLatency.WithLabelValues(info.FullMethod).Observe(time.Since(start).Seconds()) + + // Record if there was an error + if err != nil { + metricGRPCErrors.WithLabelValues(info.FullMethod).Inc() + } + }() + + // Call the handler + err = handler(srv, ss) + + return err +}