Instrumentation: log the total number of db queries per request (#54647) (#54686)

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>
(cherry picked from commit 78978048c3)

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
+}

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")
+}

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"
 )
@@ -31,6 +33,9 @@ func Logger(cfg *setting.Cfg) web.Handler {
 	return func(res http.ResponseWriter, req *http.Request, c *web.Context) {
 		start := time.Now()
 
+		// we have to init the context with the counter here to update the request
+		c.Req = c.Req.WithContext(log.InitCounter(c.Req.Context()))
+
 		rw := res.(web.ResponseWriter)
 		c.Next()
 
@@ -65,6 +70,14 @@ func Logger(cfg *setting.Cfg) web.Handler {
 				logParams = append(logParams, "traceID", traceID)
 			}
 
+			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)
+			}
+
 			if status >= 500 {
 				ctx.Logger.Error("Request Completed", logParams...)
 			} else {

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()