MemoryLimter extension as an http/grpc middleware

extension/memorylimiterextension/memorylimiter.go

@@ -5,8 +5,12 @@ package memorylimiterextension // import "go.opentelemetry.io/collector/extensio
 
 import (
 	"context"
+	"net/http"
 
 	"go.uber.org/zap"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 
 	"go.opentelemetry.io/collector/component"
 	"go.opentelemetry.io/collector/internal/memorylimiter"
@@ -38,3 +42,24 @@ func (ml *memoryLimiterExtension) Shutdown(ctx context.Context) error {
 func (ml *memoryLimiterExtension) MustRefuse() bool {
 	return ml.memLimiter.MustRefuse()
 }
+
+func (ml *memoryLimiterExtension) GetHTTPHandler(base http.Handler) (http.Handler, error) {
+	return http.HandlerFunc(func(resp http.ResponseWriter, req *http.Request) {
+		if ml.MustRefuse() {
+			http.Error(resp, http.StatusText(http.StatusTooManyRequests), http.StatusTooManyRequests)
+			return
+		}
+		base.ServeHTTP(resp, req)
+	}), nil
+}
+
+func (ml *memoryLimiterExtension) GetGRPCServerOptions() ([]grpc.ServerOption, error) {
+	return []grpc.ServerOption{grpc.UnaryInterceptor(
+		func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp any, err error) {
+			if ml.MustRefuse() {
+				return nil, status.Errorf(codes.ResourceExhausted, "RESOURCE_EXHAUSTED")
+			}
+			return handler(ctx, req)
+		},
+	)}, nil
+}