otel: use server.Config grpc server options (#4358)

korniltsev · web-flow · commit 9ea9ad385e80 · 2025-08-07T18:45:57.000+07:00
diff --git a/pkg/api/api.go b/pkg/api/api.go
@@ -9,10 +9,10 @@ import (
 	"context"
 	"flag"
 	"fmt"
-
 	"net/http"
 
 	"connectrpc.com/connect"
+
 	"github.com/felixge/fgprof"
 	"github.com/go-kit/log"
 	"github.com/grafana/dskit/kv/memberlist"
@@ -187,10 +187,10 @@ func (a *API) RegisterOverridesExporter(oe *exporter.OverridesExporter) {
 }
 
 // RegisterDistributor registers the endpoints associated with the distributor.
-func (a *API) RegisterDistributor(d *distributor.Distributor, limits *validation.Overrides, multitenancyEnabled bool) {
+func (a *API) RegisterDistributor(d *distributor.Distributor, limits *validation.Overrides, multitenancyEnabled bool, cfg server.Config) {
 	writePathOpts := a.registerOptionsWritePath(limits)
 	pyroscopeHandler := pyroscope.NewPyroscopeIngestHandler(d, a.logger)
-	otlpHandler := otlp.NewOTLPIngestHandler(d, a.logger, multitenancyEnabled)
+	otlpHandler := otlp.NewOTLPIngestHandler(cfg, d, a.logger, multitenancyEnabled)
 
 	a.RegisterRoute("/ingest", pyroscopeHandler, writePathOpts...)
 	a.RegisterRoute("/pyroscope/ingest", pyroscopeHandler, writePathOpts...)
diff --git a/pkg/ingester/otlp/ingest_handler.go b/pkg/ingester/otlp/ingest_handler.go
@@ -7,15 +7,18 @@ import (
 	"strings"
 
 	"connectrpc.com/connect"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/keepalive"
+	"google.golang.org/protobuf/proto"
+
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
 	"github.com/google/uuid"
+	"github.com/grafana/dskit/server"
 	"github.com/grafana/dskit/user"
 	pprofileotlp "go.opentelemetry.io/proto/otlp/collector/profiles/v1development"
 	v1 "go.opentelemetry.io/proto/otlp/common/v1"
-	"google.golang.org/grpc"
-	"google.golang.org/grpc/codes"
-	"google.golang.org/protobuf/proto"
 
 	"google.golang.org/grpc/status"
 
@@ -44,14 +47,14 @@ type PushService interface {
 	PushParsed(ctx context.Context, req *distirbutormodel.PushRequest) (*connect.Response[pushv1.PushResponse], error)
 }
 
-func NewOTLPIngestHandler(svc PushService, l log.Logger, me bool) Handler {
+func NewOTLPIngestHandler(cfg server.Config, svc PushService, l log.Logger, me bool) Handler {
 	h := &ingestHandler{
 		svc:                 svc,
 		log:                 l,
 		multitenancyEnabled: me,
 	}
 
-	grpcServer := grpc.NewServer()
+	grpcServer := newGrpcServer(cfg)
 	pprofileotlp.RegisterProfilesServiceServer(grpcServer, h)
 
 	h.handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -69,6 +72,34 @@ func NewOTLPIngestHandler(svc PushService, l log.Logger, me bool) Handler {
 	return h
 }
 
+func newGrpcServer(cfg server.Config) *grpc.Server {
+	grpcKeepAliveOptions := keepalive.ServerParameters{
+		MaxConnectionIdle:     cfg.GRPCServerMaxConnectionIdle,
+		MaxConnectionAge:      cfg.GRPCServerMaxConnectionAge,
+		MaxConnectionAgeGrace: cfg.GRPCServerMaxConnectionAgeGrace,
+		Time:                  cfg.GRPCServerTime,
+		Timeout:               cfg.GRPCServerTimeout,
+	}
+
+	grpcKeepAliveEnforcementPolicy := keepalive.EnforcementPolicy{
+		MinTime:             cfg.GRPCServerMinTimeBetweenPings,
+		PermitWithoutStream: cfg.GRPCServerPingWithoutStreamAllowed,
+	}
+
+	grpcOptions := []grpc.ServerOption{
+		grpc.KeepaliveParams(grpcKeepAliveOptions),
+		grpc.KeepaliveEnforcementPolicy(grpcKeepAliveEnforcementPolicy),
+		grpc.MaxRecvMsgSize(cfg.GRPCServerMaxRecvMsgSize),
+		grpc.MaxSendMsgSize(cfg.GRPCServerMaxSendMsgSize),
+		grpc.MaxConcurrentStreams(uint32(cfg.GRPCServerMaxConcurrentStreams)),
+		grpc.NumStreamWorkers(uint32(cfg.GRPCServerNumWorkers)),
+	}
+
+	grpcOptions = append(grpcOptions, cfg.GRPCOptions...)
+
+	return grpc.NewServer(grpcOptions...)
+}
+
 func (h *ingestHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	h.handler.ServeHTTP(w, r)
 }
diff --git a/pkg/ingester/otlp/ingest_handler_test.go b/pkg/ingester/otlp/ingest_handler_test.go
@@ -2,24 +2,26 @@ package otlp
 
 import (
 	"context"
+	"flag"
 	"os"
 	"sort"
 	"strings"
 	"testing"
 
-	typesv1 "github.com/grafana/pyroscope/api/gen/proto/go/types/v1"
-	"github.com/grafana/pyroscope/pkg/distributor/model"
-	phlaremodel "github.com/grafana/pyroscope/pkg/model"
-	"github.com/grafana/pyroscope/pkg/og/convert/pprof/strprofile"
-	"github.com/grafana/pyroscope/pkg/test"
-	"github.com/grafana/pyroscope/pkg/test/mocks/mockotlp"
-
+	"github.com/grafana/dskit/server"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 	v1experimental2 "go.opentelemetry.io/proto/otlp/collector/profiles/v1development"
 	v1 "go.opentelemetry.io/proto/otlp/common/v1"
 	v1experimental "go.opentelemetry.io/proto/otlp/profiles/v1development"
+
+	typesv1 "github.com/grafana/pyroscope/api/gen/proto/go/types/v1"
+	"github.com/grafana/pyroscope/pkg/distributor/model"
+	phlaremodel "github.com/grafana/pyroscope/pkg/model"
+	"github.com/grafana/pyroscope/pkg/og/convert/pprof/strprofile"
+	"github.com/grafana/pyroscope/pkg/test"
+	"github.com/grafana/pyroscope/pkg/test/mocks/mockotlp"
 )
 
 func TestGetServiceNameFromAttributes(t *testing.T) {
@@ -334,7 +336,7 @@ func TestConversion(t *testing.T) {
 						}}}}},
 				Dictionary: &b.dictionary}
 			logger := test.NewTestingLogger(t)
-			h := NewOTLPIngestHandler(svc, logger, false)
+			h := NewOTLPIngestHandler(testConfig(), svc, logger, false)
 			_, err := h.Export(context.Background(), req)
 
 			if td.expectedError == "" {
@@ -429,7 +431,7 @@ func TestSampleAttributes(t *testing.T) {
 				}}}}},
 		Dictionary: &otlpb.dictionary}
 	logger := test.NewTestingLogger(t)
-	h := NewOTLPIngestHandler(svc, logger, false)
+	h := NewOTLPIngestHandler(testConfig(), svc, logger, false)
 	_, err := h.Export(context.Background(), req)
 	assert.NoError(t, err)
 	require.Equal(t, 1, len(profiles))
@@ -593,7 +595,7 @@ func TestDifferentServiceNames(t *testing.T) {
 		Dictionary: &otlpb.dictionary}
 
 	logger := test.NewTestingLogger(t)
-	h := NewOTLPIngestHandler(svc, logger, false)
+	h := NewOTLPIngestHandler(testConfig(), svc, logger, false)
 	_, err := h.Export(context.Background(), req)
 	require.NoError(t, err)
 
@@ -649,3 +651,10 @@ func (o *otlpbuilder) addstr(s string) int32 {
 	o.dictionary.StringTable = append(o.dictionary.StringTable, s)
 	return idx
 }
+
+func testConfig() server.Config {
+	cfg := server.Config{}
+	fs := flag.NewFlagSet("test", flag.PanicOnError)
+	cfg.RegisterFlags(fs)
+	return cfg
+}
diff --git a/pkg/pyroscope/modules.go b/pkg/pyroscope/modules.go
@@ -10,6 +10,11 @@ import (
 	"time"
 
 	"connectrpc.com/connect"
+	"google.golang.org/genproto/googleapis/api/httpbody"
+	"google.golang.org/grpc/health/grpc_health_v1"
+	"google.golang.org/protobuf/encoding/protojson"
+	"gopkg.in/yaml.v3"
+
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
 	"github.com/grafana/dskit/dns"
@@ -29,10 +34,6 @@ import (
 	objstoretracing "github.com/thanos-io/objstore/tracing/opentracing"
 	"golang.org/x/net/http2"
 	"golang.org/x/net/http2/h2c"
-	"google.golang.org/genproto/googleapis/api/httpbody"
-	"google.golang.org/grpc/health/grpc_health_v1"
-	"google.golang.org/protobuf/encoding/protojson"
-	"gopkg.in/yaml.v3"
 
 	statusv1 "github.com/grafana/pyroscope/api/gen/proto/go/status/v1"
 	"github.com/grafana/pyroscope/pkg/adhocprofiles"
@@ -321,7 +322,7 @@ func (f *Pyroscope) initDistributor() (services.Service, error) {
 	if err != nil {
 		return nil, err
 	}
-	f.API.RegisterDistributor(d, f.Overrides, f.Cfg.MultitenancyEnabled)
+	f.API.RegisterDistributor(d, f.Overrides, f.Cfg.MultitenancyEnabled, f.Cfg.Server)
 	return d, nil
 }
 
