github.com/authzed/spicedb@v1.32.1-0.20240520085336-ebda56537386/internal/middleware/usagemetrics/usagemetrics.go (about)

     1  package usagemetrics
     2  
     3  import (
     4  	"context"
     5  	"strconv"
     6  	"time"
     7  
     8  	"github.com/authzed/authzed-go/pkg/responsemeta"
     9  	"github.com/authzed/grpcutil"
    10  	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors"
    11  	"github.com/prometheus/client_golang/prometheus"
    12  	"github.com/prometheus/client_golang/prometheus/promauto"
    13  	"google.golang.org/grpc"
    14  
    15  	log "github.com/authzed/spicedb/internal/logging"
    16  	dispatch "github.com/authzed/spicedb/pkg/proto/dispatch/v1"
    17  )
    18  
    19  var (
    20  	// DispatchedCountLabels are the labels that DispatchedCountHistogram will
    21  	// have have by default.
    22  	DispatchedCountLabels = []string{"method", "cached"}
    23  
    24  	// DispatchedCountHistogram is the metric that SpiceDB uses to keep track
    25  	// of the number of downstream dispatches that are performed to answer a
    26  	// single query.
    27  	DispatchedCountHistogram = promauto.NewHistogramVec(prometheus.HistogramOpts{
    28  		Namespace: "spicedb",
    29  		Subsystem: "services",
    30  		Name:      "dispatches",
    31  		Help:      "Histogram of cluster dispatches performed by the instance.",
    32  		Buckets:   []float64{1, 5, 10, 25, 50, 100, 250},
    33  	}, DispatchedCountLabels)
    34  )
    35  
    36  type reporter struct{}
    37  
    38  func (r *reporter) ServerReporter(ctx context.Context, callMeta interceptors.CallMeta) (interceptors.Reporter, context.Context) {
    39  	_, methodName := grpcutil.SplitMethodName(callMeta.FullMethod())
    40  	ctx = ContextWithHandle(ctx)
    41  	return &serverReporter{ctx: ctx, methodName: methodName}, ctx
    42  }
    43  
    44  type serverReporter struct {
    45  	interceptors.NoopReporter
    46  	ctx        context.Context
    47  	methodName string
    48  }
    49  
    50  func (r *serverReporter) PostCall(_ error, _ time.Duration) {
    51  	responseMeta := FromContext(r.ctx)
    52  	if responseMeta == nil {
    53  		responseMeta = &dispatch.ResponseMeta{}
    54  	}
    55  
    56  	err := annotateAndReportForMetadata(r.ctx, r.methodName, responseMeta)
    57  	// if context is cancelled, the stream will be closed, and gRPC will return ErrIllegalHeaderWrite
    58  	// this prevents logging unnecessary error messages
    59  	if r.ctx.Err() != nil {
    60  		return
    61  	}
    62  	if err != nil {
    63  		log.Ctx(r.ctx).Warn().Err(err).Msg("usagemetrics: could not report metadata")
    64  	}
    65  }
    66  
    67  // UnaryServerInterceptor implements a gRPC Middleware for reporting usage metrics
    68  // in both the trailer of the request, as well as to the registered prometheus
    69  // metrics.
    70  func UnaryServerInterceptor() grpc.UnaryServerInterceptor {
    71  	return interceptors.UnaryServerInterceptor(&reporter{})
    72  }
    73  
    74  // StreamServerInterceptor implements a gRPC Middleware for reporting usage metrics
    75  // in both the trailer of the request, as well as to the registered prometheus
    76  // metrics
    77  func StreamServerInterceptor() grpc.StreamServerInterceptor {
    78  	return interceptors.StreamServerInterceptor(&reporter{})
    79  }
    80  
    81  func annotateAndReportForMetadata(ctx context.Context, methodName string, metadata *dispatch.ResponseMeta) error {
    82  	DispatchedCountHistogram.WithLabelValues(methodName, "false").Observe(float64(metadata.DispatchCount))
    83  	DispatchedCountHistogram.WithLabelValues(methodName, "true").Observe(float64(metadata.CachedDispatchCount))
    84  
    85  	return responsemeta.SetResponseTrailerMetadata(ctx, map[responsemeta.ResponseMetadataTrailerKey]string{
    86  		responsemeta.DispatchedOperationsCount: strconv.Itoa(int(metadata.DispatchCount)),
    87  		responsemeta.CachedOperationsCount:     strconv.Itoa(int(metadata.CachedDispatchCount)),
    88  	})
    89  }
    90  
    91  // Create a new type to prevent context collisions
    92  type responseMetaKey string
    93  
    94  var metadataCtxKey responseMetaKey = "dispatched-response-meta"
    95  
    96  type metaHandle struct{ metadata *dispatch.ResponseMeta }
    97  
    98  // SetInContext should be called in a gRPC handler to correctly set the response metadata
    99  // for the dispatched request.
   100  func SetInContext(ctx context.Context, metadata *dispatch.ResponseMeta) {
   101  	possibleHandle := ctx.Value(metadataCtxKey)
   102  	if possibleHandle == nil {
   103  		return
   104  	}
   105  
   106  	handle := possibleHandle.(*metaHandle)
   107  	handle.metadata = metadata
   108  }
   109  
   110  // FromContext returns any metadata that was stored in the context.
   111  //
   112  // This is useful for testing that a handler is properly setting the context.
   113  func FromContext(ctx context.Context) *dispatch.ResponseMeta {
   114  	possibleHandle := ctx.Value(metadataCtxKey)
   115  	if possibleHandle == nil {
   116  		return nil
   117  	}
   118  	return possibleHandle.(*metaHandle).metadata
   119  }
   120  
   121  // ContextWithHandle creates a new context with a location to store metadata
   122  // returned from a dispatched request.
   123  //
   124  // This should only be called in middleware or testing functions.
   125  func ContextWithHandle(ctx context.Context) context.Context {
   126  	var handle metaHandle
   127  	return context.WithValue(ctx, metadataCtxKey, &handle)
   128  }