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