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

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