[DRAFT] Option to apply memory_limiter extension on all receivers

Author: dmitryax

.chloggen/timn_memorylimiterextension-HTTPServerConfig.yaml

@@ -283,21 +283,56 @@ func (gss *ServerConfig) ToListenerContext(ctx context.Context) (net.Listener, e
 	return gss.NetAddr.Listen(ctx)
 }
 
-func (gss *ServerConfig) ToServer(host component.Host, settings component.TelemetrySettings, extraOpts ...grpc.ServerOption) (*grpc.Server, error) {
-	opts, err := gss.toServerOption(host, settings)
+// ToListener returns the net.Listener constructed from the settings.
+// Deprecated: [v0.94.0] Call Listen directly on the NetAddr field.
+func (gss *ServerConfig) ToListener() (net.Listener, error) {
+	return gss.ToListenerContext(context.Background())
+}
+
+func (gss *ServerConfig) ToServer(host component.Host, settings component.TelemetrySettings, tso ...ToServerOption) (*grpc.Server, error) {
+	opts, err := gss.toServerOption(host, settings, tso)
 	if err != nil {
 		return nil, err
 	}
-	opts = append(opts, extraOpts...)
 	return grpc.NewServer(opts...), nil
 }
 
-func (gss *ServerConfig) toServerOption(host component.Host, settings component.TelemetrySettings) ([]grpc.ServerOption, error) {
+// toServerOptions has options that change the behavior of the HTTP server
+// returned by ServerConfig.ToServer().
+type toServerOptions struct {
+	memoryLimiter *component.ID
+	grpcOptions   []grpc.ServerOption
+}
+
+// ToServerOption is an option to change the behavior of the HTTP server
+// returned by ServerConfig.ToServer().
+type ToServerOption func(opts *toServerOptions)
+
+// WithGRPCServerOption adds a grpc.ServerOption to the server.
+func WithGRPCServerOption(opt grpc.ServerOption) ToServerOption {
+	return func(opts *toServerOptions) {
+		opts.grpcOptions = append(opts.grpcOptions, opt)
+	}
+}
+
+// WithMemoryLimiter sets the memory limiter to be used by the server.
+func WithMemoryLimiter(memoryLimiter *component.ID) ToServerOption {
+	return func(opts *toServerOptions) {
+		opts.memoryLimiter = memoryLimiter
+	}
+}
+
+func (gss *ServerConfig) toServerOption(host component.Host, settings component.TelemetrySettings, tso []ToServerOption) ([]grpc.ServerOption, error) {
 	switch gss.NetAddr.Transport {
 	case "tcp", "tcp4", "tcp6", "udp", "udp4", "udp6":
 		internal.WarnOnUnspecifiedHost(settings.Logger, gss.NetAddr.Endpoint)
 	}
 
+	toServerOpts := &toServerOptions{}
+	for _, o := range tso {
+		o(toServerOpts)
+	}
+
 	var opts []grpc.ServerOption
 
 	if gss.TLSSetting != nil {
@@ -369,6 +404,10 @@ func (gss *ServerConfig) toServerOption(host component.Host, settings component.
 		})
 	}
 
+	if toServerOpts.memoryLimiter != nil {
+		// TODO: Add memory limiter interceptor.
+	}
+
 	otelOpts := []otelgrpc.Option{
 		otelgrpc.WithTracerProvider(settings.TracerProvider),
 		otelgrpc.WithMeterProvider(settings.MeterProvider),
@@ -382,7 +421,7 @@ func (gss *ServerConfig) toServerOption(host component.Host, settings component.
 
 	opts = append(opts, grpc.StatsHandler(otelgrpc.NewServerHandler(otelOpts...)), grpc.ChainUnaryInterceptor(uInterceptors...), grpc.ChainStreamInterceptor(sInterceptors...))
 
-	return opts, nil
+	return append(opts, toServerOpts.grpcOptions...), nil
 }
 
 // getGRPCCompressionName returns compression name registered in grpc.

config/configgrpc/configgrpc.go

@@ -291,9 +291,6 @@ type ServerConfig struct {
 	// Additional headers attached to each HTTP response sent to the client.
 	// Header values are opaque since they may be sensitive.
 	ResponseHeaders map[string]configopaque.String `mapstructure:"response_headers"`
-
-	// MemoryLimiter is memory limiter this receiver will use to restrict incoming requests
-	MemoryLimiter *component.ID `mapstructure:"memory_limiter"`
 }
 
 // ToListener creates a net.Listener.
@@ -318,8 +315,9 @@ func (hss *ServerConfig) ToListener() (net.Listener, error) {
 // toServerOptions has options that change the behavior of the HTTP server
 // returned by ServerConfig.ToServer().
 type toServerOptions struct {
-	errHandler func(w http.ResponseWriter, r *http.Request, errorMsg string, statusCode int)
-	decoders   map[string]func(body io.ReadCloser) (io.ReadCloser, error)
+	errHandler    func(w http.ResponseWriter, r *http.Request, errorMsg string, statusCode int)
+	decoders      map[string]func(body io.ReadCloser) (io.ReadCloser, error)
+	memoryLimiter *component.ID
 }
 
 // ToServerOption is an option to change the behavior of the HTTP server
@@ -345,6 +343,12 @@ func WithDecoder(key string, dec func(body io.ReadCloser) (io.ReadCloser, error)
 	}
 }
 
+func WithMemoryLimiter(extID *component.ID) ToServerOption {
+	return func(opts *toServerOptions) {
+		opts.memoryLimiter = extID
+	}
+}
+
 // ToServer creates an http.Server from settings object.
 func (hss *ServerConfig) ToServer(host component.Host, settings component.TelemetrySettings, handler http.Handler, opts ...ToServerOption) (*http.Server, error) {
 	internal.WarnOnUnspecifiedHost(settings.Logger, hss.Endpoint)
@@ -354,8 +358,8 @@ func (hss *ServerConfig) ToServer(host component.Host, settings component.Teleme
 		o(serverOpts)
 	}
 
-	if hss.MemoryLimiter != nil {
-		ml, err := getMemoryLimiterExtension(hss.MemoryLimiter, host.GetExtensions())
+	if serverOpts.memoryLimiter != nil {
+		ml, err := getMemoryLimiterExtension(serverOpts.memoryLimiter, host.GetExtensions())
 		if err != nil {
 			return nil, err
 		}

config/confighttp/confighttp.go

@@ -51,6 +51,12 @@ type ConfigWatcher interface {
 	NotifyConfig(ctx context.Context, conf *confmap.Conf) error
 }
 
+type MemoryLimiter interface {
+	Extension
+	MustRefuse() bool
+	ApplyToAllReceivers() bool
+}
+
 // StatusWatcher is an extra interface for Extension hosted by the OpenTelemetry
 // Collector that is to be implemented by extensions interested in changes to component
 // status.

extension/extension.go

@@ -3,8 +3,9 @@
 
 package memorylimiterextension // import "go.opentelemetry.io/collector/extension/memorylimiterextension"
 
-import (
-	"go.opentelemetry.io/collector/internal/memorylimiter"
-)
+import "go.opentelemetry.io/collector/internal/memorylimiter"
 
-type Config = memorylimiter.Config
+type Config struct {
+	memorylimiter.Config `mapstructure:",squash"`
+	ApplyToAllReceivers  bool `mapstructure:"apply_to_all_receivers"`
+}

extension/memorylimiterextension/config.go

@@ -13,17 +13,21 @@ import (
 )
 
 type memoryLimiterExtension struct {
-	memLimiter *memorylimiter.MemoryLimiter
+	memLimiter          *memorylimiter.MemoryLimiter
+	applyToAllReceivers bool
 }
 
 // newMemoryLimiter returns a new memorylimiter extension.
 func newMemoryLimiter(cfg *Config, logger *zap.Logger) (*memoryLimiterExtension, error) {
-	ml, err := memorylimiter.NewMemoryLimiter(cfg, logger)
+	ml, err := memorylimiter.NewMemoryLimiter(&cfg.Config, logger)
 	if err != nil {
 		return nil, err
 	}
 
-	return &memoryLimiterExtension{memLimiter: ml}, nil
+	return &memoryLimiterExtension{
+		memLimiter:          ml,
+		applyToAllReceivers: cfg.ApplyToAllReceivers,
+	}, nil
 }
 
 func (ml *memoryLimiterExtension) Start(ctx context.Context, host component.Host) error {
@@ -38,3 +42,8 @@ func (ml *memoryLimiterExtension) Shutdown(ctx context.Context) error {
 func (ml *memoryLimiterExtension) MustRefuse() bool {
 	return ml.memLimiter.MustRefuse()
 }
+
+// ApplyToAllReceivers returns if the memory limiter must be applied to all receivers.
+func (ml *memoryLimiterExtension) ApplyToAllReceivers() bool {
+	return ml.applyToAllReceivers
+}

extension/memorylimiterextension/memorylimiter.go

@@ -44,6 +44,9 @@ type Protocols struct {
 type Config struct {
 	// Protocols is the configuration for the supported protocols, currently gRPC and HTTP (Proto and JSON).
 	Protocols `mapstructure:"protocols"`
+
+	// MemoryLimiter is memory limiter this receiver will use to restrict incoming requests to avoid OOM kills.
+	MemoryLimiter *component.ID `mapstructure:"memory_limiter"`
 }
 
 var _ component.Config = (*Config)(nil)

receiver/otlpreceiver/config.go

@@ -34,9 +34,15 @@ func NewFactory() receiver.Factory {
 		receiver.WithTraces(createTraces, metadata.TracesStability),
 		receiver.WithMetrics(createMetrics, metadata.MetricsStability),
 		receiver.WithLogs(createLog, metadata.LogsStability),
+		receiver.WithMemoryLimiter(memoryLimiter),
 	)
 }
 
+// memoryLimiter sets the memory limit for the receiver.
+func memoryLimiter(_ context.Context, _ receiver.CreateSettings, cfg component.Config) (*component.ID, error) {
+	return cfg.(*Config).MemoryLimiter, nil
+}
+
 // createDefaultConfig creates the default configuration for receiver.
 func createDefaultConfig() component.Config {
 	return &Config{

receiver/otlpreceiver/factory.go

@@ -145,7 +145,8 @@ func (r *otlpReceiver) startHTTPServer(host component.Host) error {
 	}
 
 	var err error
-	if r.serverHTTP, err = r.cfg.HTTP.ToServer(host, r.settings.TelemetrySettings, httpMux, confighttp.WithErrorHandler(errorHandler)); err != nil {
+	if r.serverHTTP, err = r.cfg.HTTP.ToServer(host, r.settings.TelemetrySettings, httpMux,
+		confighttp.WithErrorHandler(errorHandler), confighttp.WithMemoryLimiter(r.settings.MemoryLimiter)); err != nil {
 		return err
 	}
 

receiver/otlpreceiver/otlp.go

@@ -5,6 +5,7 @@ package receiver // import "go.opentelemetry.io/collector/receiver"
 
 import (
 	"context"
+	"errors"
 	"fmt"
 
 	"go.uber.org/zap"
@@ -49,6 +50,10 @@ type CreateSettings struct {
 
 	// BuildInfo can be used by components for informational purposes.
 	BuildInfo component.BuildInfo
+
+	// MemoryLimiter returns a memory limiter for the receiver.
+	// If the receiver supports memory limiting, it must use this MemoryLimiter.
+	MemoryLimiter *component.ID
 }
 
 // Factory is factory interface for receivers.
@@ -82,6 +87,12 @@ type Factory interface {
 	// LogsReceiverStability gets the stability level of the LogsReceiver.
 	LogsReceiverStability() component.StabilityLevel
 
+	// MemoryLimiter returns a memory limiter for the receiver.
+	// Implementing this method means that the receiver supports memory limiting.
+	// Even if calling this method returns nil, memory limiter can be applied implicitly, and the receiver must use
+	// MemoryLimiter provided in CreateSettings.
+	MemoryLimiter(ctx context.Context, set CreateSettings, cfg component.Config) (*component.ID, error)
+
 	unexportedFactoryFunc()
 }
 
@@ -154,6 +165,7 @@ type factory struct {
 	metricsStabilityLevel component.StabilityLevel
 	CreateLogsFunc
 	logsStabilityLevel component.StabilityLevel
+	MemoryLimiterFunc
 }
 
 func (f *factory) Type() component.Type {
@@ -198,6 +210,23 @@ func WithLogs(createLogsReceiver CreateLogsFunc, sl component.StabilityLevel) Fa
 	})
 }
 
+type MemoryLimiterFunc func(context.Context, CreateSettings, component.Config) (*component.ID, error)
+
+// MemoryLimiter implements Factory.MemoryLimiter.
+func (f MemoryLimiterFunc) MemoryLimiter(ctx context.Context, set CreateSettings, cfg component.Config) (*component.ID, error) {
+	if f == nil {
+		return nil, errors.New("the receiver doesn't support memory limiting")
+	}
+	return f(ctx, set, cfg)
+}
+
+// WithMemoryLimiter overrides the default "error not supported" implementation for MemoryLimiter.
+func WithMemoryLimiter(memoryLimiterFunc MemoryLimiterFunc) FactoryOption {
+	return factoryOptionFunc(func(o *factory) {
+		o.MemoryLimiterFunc = memoryLimiterFunc
+	})
+}
+
 // NewFactory returns a Factory.
 func NewFactory(cfgType component.Type, createDefaultConfig component.CreateDefaultConfigFunc, options ...FactoryOption) Factory {
 	f := &factory{
@@ -246,8 +275,12 @@ func (b *Builder) CreateTraces(ctx context.Context, set CreateSettings, next con
 		return nil, fmt.Errorf("receiver factory not available for: %q", set.ID)
 	}
 
+	settings, err := updateSettings(ctx, set, f, cfg)
+	if err != nil {
+		return nil, err
+	}
 	logStabilityLevel(set.Logger, f.TracesReceiverStability())
-	return f.CreateTracesReceiver(ctx, set, cfg, next)
+	return f.CreateTracesReceiver(ctx, settings, cfg, next)
 }
 
 // CreateMetrics creates a Metrics receiver based on the settings and config.
@@ -262,8 +295,12 @@ func (b *Builder) CreateMetrics(ctx context.Context, set CreateSettings, next co
 		return nil, fmt.Errorf("receiver factory not available for: %q", set.ID)
 	}
 
+	settings, err := updateSettings(ctx, set, f, cfg)
+	if err != nil {
+		return nil, err
+	}
 	logStabilityLevel(set.Logger, f.MetricsReceiverStability())
-	return f.CreateMetricsReceiver(ctx, set, cfg, next)
+	return f.CreateMetricsReceiver(ctx, settings, cfg, next)
 }
 
 // CreateLogs creates a Logs receiver based on the settings and config.
@@ -278,8 +315,12 @@ func (b *Builder) CreateLogs(ctx context.Context, set CreateSettings, next consu
 		return nil, fmt.Errorf("receiver factory not available for: %q", set.ID)
 	}
 
+	settings, err := updateSettings(ctx, set, f, cfg)
+	if err != nil {
+		return nil, err
+	}
 	logStabilityLevel(set.Logger, f.LogsReceiverStability())
-	return f.CreateLogsReceiver(ctx, set, cfg, next)
+	return f.CreateLogsReceiver(ctx, settings, cfg, next)
 }
 
 func (b *Builder) Factory(componentType component.Type) component.Factory {
@@ -296,3 +337,17 @@ func logStabilityLevel(logger *zap.Logger, sl component.StabilityLevel) {
 		logger.Info(sl.LogMessage())
 	}
 }
+
+func updateSettings(ctx context.Context, set CreateSettings, f Factory, cfg component.Config) (CreateSettings, error) {
+	ml, err := f.MemoryLimiter(ctx, set, cfg)
+
+	// Global memory limiter provided, but the receiver does not support it.
+	if set.MemoryLimiter != nil && err != nil {
+		return set, err
+	}
+
+	if ml != nil {
+		set.MemoryLimiter = ml
+	}
+	return set, nil
+}