frostfs-observability/metrics/grpc/server.go

package grpc

import (
	grpcprom "github.com/grpc-ecosystem/go-grpc-middleware/providers/prometheus"
	"github.com/prometheus/client_golang/prometheus"
	"google.golang.org/grpc"
)

// ServerMetrics is prometheus.Collector instance with server-side gRPC metrics.
type ServerMetrics interface {
	prometheus.Collector

	UnaryServerInterceptor() grpc.UnaryServerInterceptor
	StreamServerInterceptor() grpc.StreamServerInterceptor
}

// ServerOption is option to create ServerMetrics.
type ServerOption func(*serverCfg)

// NewServerMetrics returns new ServerMetrics instance.
func NewServerMetrics(opts ...ServerOption) ServerMetrics {
	cfg := &serverCfg{
		buckets: prometheus.DefBuckets,
	}
	for _, opt := range opts {
		opt(cfg)
	}
	return &serverMetrics{
		grpcprom.NewServerMetrics(
			grpcprom.WithServerHandlingTimeHistogram(
				grpcprom.WithHistogramBuckets(cfg.buckets),
			),
		),
	}
}

// WithCustomServerBuckets is ServerOption to set custom buckets for histogram metrics.
func WithCustomServerBuckets(buckets []float64) ServerOption {
	return func(cc *serverCfg) {
		cc.buckets = buckets
	}
}

type serverCfg struct {
	buckets []float64
}

type serverMetrics struct {
	*grpcprom.ServerMetrics
}

func (m *serverMetrics) UnaryServerInterceptor() grpc.UnaryServerInterceptor {
	return m.ServerMetrics.UnaryServerInterceptor()
}

func (m *serverMetrics) StreamServerInterceptor() grpc.StreamServerInterceptor {
	return m.ServerMetrics.StreamServerInterceptor()
}