From 8ffa034d2393b8948b5f5eea9a0613b15c76493e Mon Sep 17 00:00:00 2001
From: Jimmy Zelinskie <jimmy@zelinskie.com>
Date: Mon, 16 Dec 2024 18:07:22 -0500
Subject: [PATCH] pkg/cmd: add metric for logical checks

This tracks the "logical checks" which includes
not only check requests but each item in a bulk
check request.
---
 pkg/cmd/server/defaults.go | 41 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/pkg/cmd/server/defaults.go b/pkg/cmd/server/defaults.go
index a534bbf6f1..15a3a3ee83 100644
--- a/pkg/cmd/server/defaults.go
+++ b/pkg/cmd/server/defaults.go
@@ -22,6 +22,7 @@ import (
 	"github.com/jzelinskie/cobrautil/v2/cobraproclimits"
 	"github.com/jzelinskie/cobrautil/v2/cobrazerolog"
 	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"github.com/rs/zerolog"
 	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
@@ -30,6 +31,8 @@ import (
 	"google.golang.org/grpc/codes"
 
 	"github.com/authzed/authzed-go/pkg/requestmeta"
+	v1 "github.com/authzed/authzed-go/proto/authzed/api/v1"
+	"github.com/authzed/grpcutil"
 
 	"github.com/authzed/spicedb/internal/dispatch"
 	"github.com/authzed/spicedb/internal/logging"
@@ -170,6 +173,7 @@ const (
 	DefaultMiddlewareGRPCAuth      = "grpcauth"
 	DefaultMiddlewareGRPCProm      = "grpcprom"
 	DefaultMiddlewareServerVersion = "serverversion"
+	DefaultMiddlewareLogicalChecks = "logicalchecks"
 
 	DefaultInternalMiddlewareDispatch       = "dispatch"
 	DefaultInternalMiddlewareDatastore      = "datastore"
@@ -333,6 +337,11 @@ func DefaultUnaryMiddleware(opts MiddlewareOption) (*MiddlewareChain[grpc.UnaryS
 			WithInterceptor(serverversion.UnaryServerInterceptor(opts.EnableVersionResponse)).
 			Done(),
 
+		NewUnaryMiddleware().
+			WithName(DefaultMiddlewareLogicalChecks).
+			WithInterceptor(LogicalChecksMetricUnary).
+			Done(),
+
 		NewUnaryMiddleware().
 			WithName(DefaultInternalMiddlewareDispatch).
 			WithInternal(true).
@@ -406,6 +415,11 @@ func DefaultStreamingMiddleware(opts MiddlewareOption) (*MiddlewareChain[grpc.St
 			WithInterceptor(serverversion.StreamServerInterceptor(opts.EnableVersionResponse)).
 			Done(),
 
+		NewStreamMiddleware().
+			WithName(DefaultMiddlewareLogicalChecks).
+			WithInterceptor(noopStreamInterceptor).
+			Done(),
+
 		NewStreamMiddleware().
 			WithName(DefaultInternalMiddlewareDispatch).
 			WithInternal(true).
@@ -428,6 +442,10 @@ func DefaultStreamingMiddleware(opts MiddlewareOption) (*MiddlewareChain[grpc.St
 	return &chain, err
 }
 
+func noopStreamInterceptor(srv any, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
+	return handler(srv, stream)
+}
+
 func determineEventsToLog(opts MiddlewareOption) grpclog.Option {
 	eventsToLog := []grpclog.LoggableEvent{grpclog.FinishCall}
 	if opts.EnableRequestLog {
@@ -465,6 +483,29 @@ func DefaultDispatchMiddleware(logger zerolog.Logger, authFunc grpcauth.AuthFunc
 		}
 }
 
+var LogicalChecksCounter = promauto.NewCounterVec(prometheus.CounterOpts{
+	Namespace: "spicedb",
+	Subsystem: "logical",
+	Name:      "permission_checks_total",
+	Help:      "Count of the checks across individual and bulk requests",
+}, []string{"grpc_method", "grpc_service", "grpc_type"})
+
+func LogicalChecksMetricUnary(ctx context.Context, request any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
+	response, err := handler(ctx, request)
+	if err == nil {
+		switch m := request.(type) {
+		case *v1.CheckPermissionRequest:
+			svc, method := grpcutil.SplitMethodName(info.FullMethod)
+			LogicalChecksCounter.WithLabelValues(method, svc, "unary").Add(1)
+		case *v1.CheckBulkPermissionsRequest:
+			svc, method := grpcutil.SplitMethodName(info.FullMethod)
+			LogicalChecksCounter.WithLabelValues(method, svc, "unary").Add(float64(len(m.GetItems())))
+		default:
+		}
+	}
+	return response, err
+}
+
 func InterceptorLogger(l zerolog.Logger) grpclog.Logger {
 	return grpclog.LoggerFunc(func(ctx context.Context, lvl grpclog.Level, msg string, fields ...any) {
 		l := l.With().Fields(fields).Logger()