telemetry: Refactor GRPC Server Handler (#756)

* Refactor GRPC server stats handler location
This commit is contained in:
Travis Groth 2020-05-22 13:36:55 -04:00 committed by GitHub
parent e2a7149c36
commit ca5f68e371
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 106 additions and 21 deletions

View file

@ -4,12 +4,13 @@ import (
"context"
"strings"
"github.com/pomerium/pomerium/internal/log"
"go.opencensus.io/plugin/ocgrpc"
"go.opencensus.io/stats/view"
"go.opencensus.io/tag"
"google.golang.org/grpc"
grpcstats "google.golang.org/grpc/stats"
"github.com/pomerium/pomerium/internal/log"
)
// GRPC Views
@ -110,6 +111,9 @@ var (
// GRPCClientInterceptor creates a UnaryClientInterceptor which updates the RPC
// context with metric tag metadata
//
// TODO: This handler will NOT currently propagate B3 headers to upstream servers. See
// GRPCServerStatsHandler for changes required
func GRPCClientInterceptor(service string) grpc.UnaryClientInterceptor {
return func(
ctx context.Context,
@ -147,17 +151,18 @@ func GRPCClientInterceptor(service string) grpc.UnaryClientInterceptor {
}
// GRPCServerStatsHandler provides a grpc stats.Handler for a pomerium service to add tags and track
// metrics to server side calls
type GRPCServerStatsHandler struct {
// GRPCServerMetricsHandler implements a telemetry tagRPCHandler methods for metrics
type GRPCServerMetricsHandler struct {
service string
grpcstats.Handler
}
// TagRPC implements grpc.stats.Handler and adds tags to the context of a given RPC
func (h *GRPCServerStatsHandler) TagRPC(ctx context.Context, tagInfo *grpcstats.RPCTagInfo) context.Context {
// NewGRPCServerMetricsHandler creates a new GRPCServerStatsHandler for a pomerium service
func NewGRPCServerMetricsHandler(service string) *GRPCServerMetricsHandler {
return &GRPCServerMetricsHandler{service: service}
}
handledCtx := h.Handler.TagRPC(ctx, tagInfo)
// TagRPC handles adding any metrics related values to the incoming context
func (h *GRPCServerMetricsHandler) TagRPC(ctx context.Context, tagInfo *grpcstats.RPCTagInfo) context.Context {
// Split the method into parts for better slicing
rpcInfo := strings.SplitN(tagInfo.FullMethodName, "/", 3)
@ -169,21 +174,16 @@ func (h *GRPCServerStatsHandler) TagRPC(ctx context.Context, tagInfo *grpcstats.
}
taggedCtx, tagErr := tag.New(
handledCtx,
ctx,
tag.Upsert(TagKeyService, h.service),
tag.Upsert(TagKeyGRPCMethod, rpcMethod),
tag.Upsert(TagKeyGRPCService, rpcService),
)
if tagErr != nil {
log.Warn().Err(tagErr).Str("context", "GRPCServerStatsHandler").Msg("telemetry/metrics: failed to create context")
return handledCtx
return ctx
}
return taggedCtx
}
// NewGRPCServerStatsHandler creates a new GRPCServerStatsHandler for a pomerium service
func NewGRPCServerStatsHandler(service string) grpcstats.Handler {
return &GRPCServerStatsHandler{service: service, Handler: &ocgrpc.ServerHandler{}}
}

View file

@ -11,6 +11,8 @@ import (
"google.golang.org/grpc/status"
)
var statsHandler = &ocgrpc.ServerHandler{}
type testProto struct {
message string
}
@ -118,15 +120,17 @@ func Test_GRPCClientInterceptor(t *testing.T) {
}
}
func mockServerRPCHandle(statsHandler stats.Handler, method string, errorCode error) {
func mockServerRPCHandle(metricsHandler *GRPCServerMetricsHandler, method string, errorCode error) {
message := "hello"
ctx := statsHandler.TagRPC(context.Background(), &stats.RPCTagInfo{FullMethodName: method})
ctx = metricsHandler.TagRPC(ctx, &stats.RPCTagInfo{FullMethodName: method})
statsHandler.HandleRPC(ctx, &stats.InPayload{Client: false, Length: len(message)})
statsHandler.HandleRPC(ctx, &stats.OutPayload{Client: false, Length: len(message)})
statsHandler.HandleRPC(ctx, &stats.End{Client: false, Error: errorCode})
}
func Test_GRPCServerStatsHandler(t *testing.T) {
func Test_GRPCServerMetricsHandler(t *testing.T) {
tests := []struct {
name string
method string
@ -171,8 +175,8 @@ func Test_GRPCServerStatsHandler(t *testing.T) {
view.Unregister(GRPCServerViews...)
view.Register(GRPCServerViews...)
statsHandler := NewGRPCServerStatsHandler("test_service")
mockServerRPCHandle(statsHandler, tt.method, tt.errorCode)
metricsHandler := NewGRPCServerMetricsHandler("test_service")
mockServerRPCHandle(metricsHandler, tt.method, tt.errorCode)
testDataRetrieval(GRPCServerResponseSizeView, t, tt.wantgrpcServerResponseSize)
testDataRetrieval(GRPCServerRequestDurationView, t, tt.wantgrpcServerRequestDuration)