From 78978048c3ea7de389a25c8c111c74923d103fab Mon Sep 17 00:00:00 2001
From: Carl Bergquist <carl.bergquist@gmail.com>
Date: Mon, 5 Sep 2022 07:39:22 +0200
Subject: [PATCH] Instrumentation: log the total number of db queries per
 request (#54647)

Signed-off-by: bergquist <carl.bergquist@gmail.com>
Co-authored-by: Emil Tullstedt <emil.tullstedt@grafana.com>
Co-authored-by: Dave Henderson <dave.henderson@grafana.com>
---
 pkg/infra/log/databaseCounter.go          | 44 +++++++++++++++++++++++
 pkg/infra/log/databaseCounter_test.go     | 19 ++++++++++
 pkg/middleware/logger.go                  | 13 +++++++
 pkg/services/sqlstore/database_wrapper.go |  2 ++
 4 files changed, 78 insertions(+)
 create mode 100644 pkg/infra/log/databaseCounter.go
 create mode 100644 pkg/infra/log/databaseCounter_test.go

diff --git a/pkg/infra/log/databaseCounter.go b/pkg/infra/log/databaseCounter.go
new file mode 100644
index 00000000000..8043f37e1a2
--- /dev/null
+++ b/pkg/infra/log/databaseCounter.go
@@ -0,0 +1,44 @@
+package log
+
+import (
+	"context"
+	"sync/atomic"
+)
+
+type contextKey struct{}
+
+var dbCallCounterNameKey = contextKey{}
+
+// InitCounter creates a pointer on the context that can be incremented later
+func InitCounter(ctx context.Context) context.Context {
+	var ptr *int64 = new(int64)
+	return context.WithValue(ctx, dbCallCounterNameKey, ptr)
+}
+
+// IncDBCallCounter increments the database counter on the context.
+func IncDBCallCounter(ctx context.Context) context.Context {
+	if val := ctx.Value(dbCallCounterNameKey); val == nil {
+		ctx = InitCounter(ctx)
+	}
+
+	if val := ctx.Value(dbCallCounterNameKey); val != nil {
+		v2, ok := val.(*int64)
+		if ok {
+			atomic.AddInt64(v2, 1)
+		}
+	}
+
+	return ctx
+}
+
+// TotalDBCallCount returns the total number of requests for the context
+func TotalDBCallCount(ctx context.Context) int64 {
+	if val := ctx.Value(dbCallCounterNameKey); val != nil {
+		v2, ok := val.(*int64)
+		if ok {
+			return *v2
+		}
+	}
+
+	return 0
+}
diff --git a/pkg/infra/log/databaseCounter_test.go b/pkg/infra/log/databaseCounter_test.go
new file mode 100644
index 00000000000..18e894f962e
--- /dev/null
+++ b/pkg/infra/log/databaseCounter_test.go
@@ -0,0 +1,19 @@
+package log
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCountingDatabaseCalls(t *testing.T) {
+	ctx := context.Background()
+
+	ctx = IncDBCallCounter(ctx)
+	ctx = IncDBCallCounter(ctx)
+	ctx = IncDBCallCounter(ctx)
+
+	count := TotalDBCallCount(ctx)
+	assert.Equal(t, int64(3), count, "expect counter to increase three times")
+}
diff --git a/pkg/middleware/logger.go b/pkg/middleware/logger.go
index 67bbf95ce09..06abe72bcb5 100644
--- a/pkg/middleware/logger.go
+++ b/pkg/middleware/logger.go
@@ -20,9 +20,11 @@ import (
 	"net/url"
 	"time"
 
+	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/infra/tracing"
 	"github.com/grafana/grafana/pkg/models"
 	"github.com/grafana/grafana/pkg/services/contexthandler"
+	"github.com/grafana/grafana/pkg/services/featuremgmt"
 	"github.com/grafana/grafana/pkg/setting"
 	"github.com/grafana/grafana/pkg/web"
 )
@@ -32,6 +34,9 @@ func Logger(cfg *setting.Cfg) web.Middleware {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			start := time.Now()
 
+			// we have to init the context with the counter here to update the request
+			r = r.WithContext(log.InitCounter(r.Context()))
+
 			rw := web.Rw(w, r)
 			next.ServeHTTP(rw, r)
 
@@ -61,6 +66,14 @@ func Logger(cfg *setting.Cfg) web.Middleware {
 					"referer", SanitizeURL(ctx, r.Referer()),
 				}
 
+				if cfg.IsFeatureToggleEnabled(featuremgmt.FlagDatabaseMetrics) {
+					logParams = append(logParams, "db_call_count", log.TotalDBCallCount(ctx.Req.Context()))
+				}
+
+				if handler, exist := routeOperationName(ctx.Req); exist {
+					logParams = append(logParams, "handler", handler)
+				}
+
 				traceID := tracing.TraceIDFromContext(ctx.Req.Context(), false)
 				if traceID != "" {
 					logParams = append(logParams, "traceID", traceID)
diff --git a/pkg/services/sqlstore/database_wrapper.go b/pkg/services/sqlstore/database_wrapper.go
index 400f823a55e..caf5dee5564 100644
--- a/pkg/services/sqlstore/database_wrapper.go
+++ b/pkg/services/sqlstore/database_wrapper.go
@@ -93,6 +93,8 @@ func (h *databaseQueryWrapper) instrument(ctx context.Context, status string, qu
 		histogram.Observe(elapsed.Seconds())
 	}
 
+	ctx = log.IncDBCallCounter(ctx)
+
 	_, span := h.tracer.Start(ctx, "database query")
 	defer span.End()