(configgrpc): configure memorylimiterextension when setting up grpc servers

.chloggen/configgrpc-configure-memory-limiter-extension.yaml

@@ -0,0 +1,25 @@
+# Use this changelog template to create an entry for release notes.
+
+# One of 'breaking', 'deprecation', 'new_component', 'enhancement', 'bug_fix'
+change_type: enhancement
+
+# The name of the component, or a single word describing the area of concern, (e.g. otlpreceiver)
+component: configgrpc
+
+# A brief description of the change.  Surround your text with quotes ("") if it needs to start with a backtick (`).
+note: Adds option to configure a memory limiter extension within grpc servers.
+
+# One or more tracking issues or pull requests related to the change
+issues: [9591]
+
+# (Optional) One or more lines of additional information to render under the primary note.
+# These lines will be padded with 2 spaces and then inserted directly into the document.
+# Use pipe (|) for multiline entries.
+subtext:
+
+# Optional: The change log or logs in which this entry should be included.
+# e.g. '[user]' or '[user, api]'
+# Include 'user' if the change is relevant to end users.
+# Include 'api' if there is a change to a library API.
+# Default: '[user]'
+change_logs: []

config/configgrpc/README.md

@@ -27,6 +27,7 @@ README](../configtls/README.md).
 - [`read_buffer_size`](https://godoc.org/google.golang.org/grpc#ReadBufferSize)
 - [`write_buffer_size`](https://godoc.org/google.golang.org/grpc#WriteBufferSize)
 - [`auth`](../configauth/README.md)
+- [`memory_limiter`](../../extension/memorylimiterextension/README.md) rejects incoming requests once the memory utilization grows above configured limits.
 
 Please note that [`per_rpc_auth`](https://pkg.go.dev/google.golang.org/grpc#PerRPCCredentials) which allows the credentials to send for every RPC is now moved to become an [extension](https://github.com/open-telemetry/opentelemetry-collector-contrib/blob/main/extension/bearertokenauthextension). Note that this feature isn't about sending the headers only during the initial connection as an `authorization` header under the `headers` would do: this is sent for every RPC performed during an established connection.
 

config/configgrpc/configgrpc.go

@@ -17,12 +17,14 @@ import (
 	"go.opentelemetry.io/otel"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/encoding/gzip"
 	"google.golang.org/grpc/keepalive"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/peer"
+	"google.golang.org/grpc/status"
 
 	"go.opentelemetry.io/collector/client"
 	"go.opentelemetry.io/collector/component"
@@ -35,7 +37,10 @@ import (
 	"go.opentelemetry.io/collector/extension/auth"
 )
 
-var errMetadataNotFound = errors.New("no request metadata found")
+var (
+	errMetadataNotFound   = errors.New("no request metadata found")
+	errMemoryLimitReached = status.Error(codes.ResourceExhausted, "Memory limit has been reached, too many requests.")
+)
 
 // KeepaliveClientConfig exposes the keepalive.ClientParameters to be used by the exporter.
 // Refer to the original data-structure for the meaning of each parameter:
@@ -88,6 +93,9 @@ type ClientConfig struct {
 
 	// Auth configuration for outgoing RPCs.
 	Auth *configauth.Authentication `mapstructure:"auth"`
+
+	// MemoryLimiter is memory limiter this receiver will use to restrict incoming requests
+	MemoryLimiter *component.ID `mapstructure:"memory_limiter"`
 }
 
 // KeepaliveServerConfig is the configuration for keepalive.
@@ -148,8 +156,12 @@ type ServerConfig struct {
 	// Include propagates the incoming connection's metadata to downstream consumers.
 	// Experimental: *NOTE* this option is subject to change or removal in the future.
 	IncludeMetadata bool `mapstructure:"include_metadata"`
+
+	MemoryLimiter *component.ID `mapstructure:"memory_limiter"`
 }
 
+type memoryLimiterExtension = interface{ MustRefuse() bool }
+
 // SanitizedEndpoint strips the prefix of either http:// or https:// from configgrpc.ClientConfig.Endpoint.
 func (gcs *ClientConfig) SanitizedEndpoint() string {
 	switch {
@@ -360,6 +372,20 @@ func (gss *ServerConfig) toServerOption(host component.Host, settings component.
 		})
 	}
 
+	if gss.MemoryLimiter != nil {
+		memoryLimiter, err := getMemoryLimiterExtension(gss.MemoryLimiter, host.GetExtensions())
+		if err != nil {
+			return nil, err
+		}
+
+		uInterceptors = append(uInterceptors, func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp any, err error) {
+			return memoryLimiterUnaryServerInterceptor(ctx, req, info, handler, memoryLimiter)
+		})
+		sInterceptors = append(sInterceptors, func(srv any, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
+			return memoryLimiterStreamServerInterceptor(srv, ss, info, handler, memoryLimiter)
+		})
+	}
+
 	otelOpts := []otelgrpc.Option{
 		otelgrpc.WithTracerProvider(settings.TracerProvider),
 		otelgrpc.WithMeterProvider(settings.MeterProvider),
@@ -376,6 +402,34 @@ func (gss *ServerConfig) toServerOption(host component.Host, settings component.
 	return opts, nil
 }
 
+func memoryLimiterUnaryServerInterceptor(ctx context.Context, req any, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler, ml memoryLimiterExtension) (any, error) {
+	if ml.MustRefuse() {
+		return nil, errMemoryLimitReached
+	}
+
+	return handler(ctx, req)
+}
+
+func memoryLimiterStreamServerInterceptor(srv any, stream grpc.ServerStream, _ *grpc.StreamServerInfo, handler grpc.StreamHandler, ml memoryLimiterExtension) error {
+	ctx := stream.Context()
+	if ml.MustRefuse() {
+		return errMemoryLimitReached
+	}
+
+	return handler(srv, wrapServerStream(ctx, stream))
+}
+
+func getMemoryLimiterExtension(extID *component.ID, extensions map[component.ID]component.Component) (memoryLimiterExtension, error) {
+	if ext, found := extensions[*extID]; found {
+		if server, ok := ext.(memoryLimiterExtension); ok {
+			return server, nil
+		}
+		return nil, fmt.Errorf("requested MemoryLimiter, %s, is not a memoryLimiterExtension", extID)
+	}
+
+	return nil, fmt.Errorf("failed to resolve memoryLimiterExtension %q: %s", extID, "memory limiter extension not found")
+}
+
 // getGRPCCompressionName returns compression name registered in grpc.
 func getGRPCCompressionName(compressionType configcompression.Type) (string, error) {
 	switch compressionType {

config/configgrpc/configgrpc_test.go

@@ -6,6 +6,7 @@ package configgrpc
 import (
 	"context"
 	"errors"
+	"fmt"
 	"net"
 	"os"
 	"path/filepath"
@@ -32,6 +33,7 @@ import (
 	"go.opentelemetry.io/collector/config/configtls"
 	"go.opentelemetry.io/collector/extension/auth"
 	"go.opentelemetry.io/collector/extension/auth/authtest"
+	"go.opentelemetry.io/collector/extension/extensiontest"
 	"go.opentelemetry.io/collector/pdata/ptrace/ptraceotlp"
 )
 
@@ -261,6 +263,225 @@ func TestGrpcServerAuthSettings_Deprecated(t *testing.T) {
 	assert.NotNil(t, srv)
 }
 
+type mockMemoryLimiterExtension struct {
+	iD         component.ID
+	mustRefuse bool
+	component.StartFunc
+	component.ShutdownFunc
+}
+
+func (mml *mockMemoryLimiterExtension) MustRefuse() bool {
+	return mml.mustRefuse
+}
+
+func TestGetMemoryLimiterExtension(t *testing.T) {
+	badID := component.NewID("badmemlimiter")
+	notMLExtensionErr := fmt.Errorf("requested MemoryLimiter, %s, is not a memoryLimiterExtension", badID)
+
+	missingID := component.NewID("missingmemlimiter")
+	missingMLExtensionErr := fmt.Errorf("failed to resolve memoryLimiterExtension %q: %s", missingID, "memory limiter extension not found")
+
+	validID := component.NewID("memorylimiter")
+
+	tests := []struct {
+		name        string
+		componentID component.ID
+		// getExtErr refers to whether the requested memory limiter extension was successfully found.
+		getExtErr error
+	}{
+		{
+			name:        "memory extension found",
+			componentID: validID,
+			getExtErr:   nil,
+		},
+		{
+			name:        "not a memory limiter extension",
+			componentID: badID,
+			getExtErr:   notMLExtensionErr,
+		},
+		{
+			name:        "memory limiter extension not found",
+			componentID: missingID,
+			getExtErr:   missingMLExtensionErr,
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			comID := test.componentID
+			gss := &ServerConfig{
+				NetAddr: confignet.AddrConfig{
+					Endpoint:  "localhost:0",
+					Transport: "tcp",
+				},
+				MemoryLimiter: &comID,
+			}
+
+			nopExt, err := extensiontest.NewNopFactory().CreateExtension(context.Background(), extensiontest.NewNopCreateSettings(), ServerConfig{})
+			assert.NoError(t, err)
+			ml := &mockMemoryLimiterExtension{iD: comID}
+			extList := map[component.ID]component.Component{
+				component.NewID("memorylimiter"): ml,
+				badID:                            nopExt,
+			}
+
+			host := &mockHost{
+				ext: extList,
+			}
+
+			// ToServer calls getMemoryLimiterExtension().
+			srv, err := gss.ToServer(host, componenttest.NewNopTelemetrySettings())
+
+			// desired extension was not found.
+			if test.getExtErr != nil {
+				assert.Equal(t, test.getExtErr.Error(), err.Error())
+				assert.Nil(t, srv)
+				return
+			}
+
+			assert.NoError(t, err)
+		})
+	}
+
+}
+
+func TestGrpcUnaryServerMemoryLimiterSettings(t *testing.T) {
+	tests := []struct {
+		name    string
+		refused bool
+		// interceptErr refers to whether memorylimiterextension allowed the request.
+		interceptErr error
+	}{
+		{
+			name:         "unary memory limiter extension accept",
+			refused:      false,
+			interceptErr: nil,
+		},
+		{
+			name:         "unary memory limiter extension refuse",
+			refused:      true,
+			interceptErr: errMemoryLimitReached,
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			comID := component.NewID("memorylimiter")
+			gss := &ServerConfig{
+				NetAddr: confignet.AddrConfig{
+					Endpoint:  "localhost:0",
+					Transport: "tcp",
+				},
+				MemoryLimiter: &comID,
+			}
+
+			ml := &mockMemoryLimiterExtension{iD: comID, mustRefuse: test.refused}
+			extList := map[component.ID]component.Component{
+				comID: ml,
+			}
+
+			host := &mockHost{
+				ext: extList,
+			}
+
+			srv, err := gss.ToServer(host, componenttest.NewNopTelemetrySettings())
+
+			// found extension so finish setting up server and client to test interceptor.
+			assert.NoError(t, err)
+			mock := &grpcTraceServer{}
+
+			ptraceotlp.RegisterGRPCServer(srv, mock)
+
+			defer srv.Stop()
+
+			l, err := gss.NetAddr.Listen(context.Background())
+			require.NoError(t, err)
+
+			go func() {
+				_ = srv.Serve(l)
+			}()
+
+			// setup client
+			gcs := &ClientConfig{
+				Endpoint: l.Addr().String(),
+				TLSSetting: configtls.ClientConfig{
+					Insecure: true,
+				},
+			}
+
+			tt, err := componenttest.SetupTelemetry(comID)
+			require.NoError(t, err)
+			defer func() {
+				require.NoError(t, tt.Shutdown(context.Background()))
+			}()
+
+			grpcClientConn, errClient := gcs.ToClientConn(context.Background(), componenttest.NewNopHost(), tt.TelemetrySettings())
+			require.NoError(t, errClient)
+			defer func() { assert.NoError(t, grpcClientConn.Close()) }()
+
+			cl := ptraceotlp.NewGRPCClient(grpcClientConn)
+			ctx, cancelFunc := context.WithTimeout(context.Background(), 2*time.Second)
+			defer cancelFunc()
+
+			resp, errResp := cl.Export(ctx, ptraceotlp.NewExportRequest())
+			assert.ErrorIs(t, test.interceptErr, errResp)
+
+			if test.interceptErr != nil {
+				assert.Equal(t, resp, ptraceotlp.ExportResponse{})
+			} else {
+				assert.NotNil(t, resp)
+				assert.NotEqual(t, resp, ptraceotlp.ExportResponse{})
+			}
+		})
+	}
+}
+
+func TestGrpcStreamServerMemoryLimiterSettings(t *testing.T) {
+	tests := []struct {
+		name    string
+		refused bool
+		// interceptErr refers to whether memorylimiterextension allowed the request.
+		interceptErr error
+	}{
+		{
+			name:         "stream memory limiter extension accept",
+			refused:      false,
+			interceptErr: nil,
+		},
+		{
+			name:         "stream memory limiter extension refuse",
+			refused:      true,
+			interceptErr: errMemoryLimitReached,
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			comID := component.NewID("memorylimiter")
+
+			ml := &mockMemoryLimiterExtension{iD: comID, mustRefuse: test.refused}
+			handlerCalled := false
+			handler := func(_ any, _ grpc.ServerStream) error {
+				handlerCalled = true
+				return nil
+			}
+			ctx := metadata.NewIncomingContext(context.Background(), metadata.Pairs("authorization", "some-auth-data"))
+			streamServer := &mockServerStream{
+				ctx: ctx,
+			}
+
+			// test
+			err := memoryLimiterStreamServerInterceptor(nil, streamServer, &grpc.StreamServerInfo{}, handler, ml)
+
+			// verify
+			assert.ErrorIs(t, test.interceptErr, err)
+			if test.refused {
+				assert.False(t, handlerCalled)
+			} else {
+				assert.True(t, handlerCalled)
+			}
+		})
+	}
+
+}
+
 func TestGRPCClientSettingsError(t *testing.T) {
 	tt, err := componenttest.SetupTelemetry(componentID)
 	require.NoError(t, err)

config/configgrpc/go.mod

@@ -13,6 +13,7 @@ require (
 	go.opentelemetry.io/collector/config/configopaque v1.3.0
 	go.opentelemetry.io/collector/config/configtls v0.96.0
 	go.opentelemetry.io/collector/config/internal v0.96.0
+	go.opentelemetry.io/collector/extension v0.96.0
 	go.opentelemetry.io/collector/extension/auth v0.96.0
 	go.opentelemetry.io/collector/pdata v1.3.0
 	go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.47.0
@@ -51,7 +52,6 @@ require (
 	github.com/prometheus/procfs v0.12.0 // indirect
 	go.opentelemetry.io/collector/config/configtelemetry v0.96.0 // indirect
 	go.opentelemetry.io/collector/confmap v0.96.0 // indirect
-	go.opentelemetry.io/collector/extension v0.96.0 // indirect
 	go.opentelemetry.io/collector/featuregate v1.3.0 // indirect
 	go.opentelemetry.io/otel/exporters/prometheus v0.46.0 // indirect
 	go.opentelemetry.io/otel/metric v1.24.0 // indirect