github.com/weaveworks/common@v0.0.0-20230728070032-dd9e68f319d5/middleware/grpc_stats.go (about)

     1  package middleware
     2  
     3  import (
     4  	"context"
     5  
     6  	"github.com/prometheus/client_golang/prometheus"
     7  	"google.golang.org/grpc/stats"
     8  )
     9  
    10  // NewStatsHandler creates handler that can be added to gRPC server options to track received and sent message sizes.
    11  func NewStatsHandler(receivedPayloadSize, sentPayloadSize *prometheus.HistogramVec, inflightRequests *prometheus.GaugeVec) stats.Handler {
    12  	return &grpcStatsHandler{
    13  		receivedPayloadSize: receivedPayloadSize,
    14  		sentPayloadSize:     sentPayloadSize,
    15  		inflightRequests:    inflightRequests,
    16  	}
    17  }
    18  
    19  type grpcStatsHandler struct {
    20  	receivedPayloadSize *prometheus.HistogramVec
    21  	sentPayloadSize     *prometheus.HistogramVec
    22  	inflightRequests    *prometheus.GaugeVec
    23  }
    24  
    25  // Custom type to hide it from other packages.
    26  type contextKey int
    27  
    28  const (
    29  	contextKeyMethodName contextKey = 1
    30  )
    31  
    32  func (g *grpcStatsHandler) TagRPC(ctx context.Context, info *stats.RPCTagInfo) context.Context {
    33  	return context.WithValue(ctx, contextKeyMethodName, info.FullMethodName)
    34  }
    35  
    36  func (g *grpcStatsHandler) HandleRPC(ctx context.Context, rpcStats stats.RPCStats) {
    37  	// We use full method name from context, because not all RPCStats structs have it.
    38  	fullMethodName, ok := ctx.Value(contextKeyMethodName).(string)
    39  	if !ok {
    40  		return
    41  	}
    42  
    43  	switch s := rpcStats.(type) {
    44  	case *stats.Begin:
    45  		g.inflightRequests.WithLabelValues(gRPC, fullMethodName).Inc()
    46  	case *stats.End:
    47  		g.inflightRequests.WithLabelValues(gRPC, fullMethodName).Dec()
    48  	case *stats.InHeader:
    49  		// Ignore incoming headers.
    50  	case *stats.InPayload:
    51  		g.receivedPayloadSize.WithLabelValues(gRPC, fullMethodName).Observe(float64(s.WireLength))
    52  	case *stats.InTrailer:
    53  		// Ignore incoming trailers.
    54  	case *stats.OutHeader:
    55  		// Ignore outgoing headers.
    56  	case *stats.OutPayload:
    57  		g.sentPayloadSize.WithLabelValues(gRPC, fullMethodName).Observe(float64(s.WireLength))
    58  	case *stats.OutTrailer:
    59  		// Ignore outgoing trailers. OutTrailer doesn't have valid WireLength (there is a deprecated field, always set to 0).
    60  	}
    61  }
    62  
    63  func (g *grpcStatsHandler) TagConn(ctx context.Context, _ *stats.ConnTagInfo) context.Context {
    64  	return ctx
    65  }
    66  
    67  func (g *grpcStatsHandler) HandleConn(_ context.Context, _ stats.ConnStats) {
    68  	// Not interested.
    69  }