MemoryLimiter extension as an http/grpc middleware

extension/memorylimiterextension/README.md

@@ -14,6 +14,27 @@ The memory limiter extension is used to prevent out of memory situations on
 the collector. The extension will potentially replace the Memory Limiter Processor. 
 It provides better guarantees from running out of memory as it will be used by the 
 receivers to reject requests before converting them into OTLP. All the configurations 
-are the same as Memory Limiter Processor. The extension is under development and does nothing.
+are the same as Memory Limiter Processor. 
+
+## How to use extension as a middleware/interceptor 
+This extension can be used as an extension to otlp receiver.
+
+```
+receivers:
+  otlp:
+    protocols:
+      grpc:
+        middlewares:
+          - id: memory_limiter
+      http:
+        middlewares:
+          - id: memory_limiter
+
+extensions:
+  memory_limiter:
+    check_interval: 1s
+    limit_percentage: 1
+    spike_limit_percentage: 0.05
+```
 
 see [memorylimiterprocessor](../../processor/memorylimiterprocessor/README.md) for additional details

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
+}

extension/memorylimiterextension/memorylimiter_test.go

@@ -14,6 +14,7 @@ import (
 	"go.uber.org/zap"
 
 	"go.opentelemetry.io/collector/component/componenttest"
+	"go.opentelemetry.io/collector/extension/extensionmiddleware"
 	"go.opentelemetry.io/collector/internal/memorylimiter"
 	"go.opentelemetry.io/collector/internal/memorylimiter/iruntime"
 )
@@ -95,3 +96,8 @@ func TestMemoryPressureResponse(t *testing.T) {
 		})
 	}
 }
+
+func TestMiddleware(t *testing.T) {
+	var _ extensionmiddleware.HTTPServer = memoryLimiterExtension{}
+	var _ extensionmiddleware.GRPCServer = memoryLimiterExtension{}
+}