ratelimitprocessor: Wrap error with code for 429 (#656)

Wraps the error returned by the rate limiter with the grpc code library
to ensure that the code is propagated to the client. For non 5xx errors,
the code is set to `codes.ResourceExhausted`, which is the appropriate
code for rate limiting errors.

Part of elastic/hosted-otel-collector#991

Signed-off-by: Marc Lopez Rubio <[email protected]>

Author: marclop

processor/ratelimitprocessor/gubernator.go

@@ -29,7 +29,9 @@ import (
 
 	"github.com/sirupsen/logrus"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/status"
 
 	"github.com/gubernator-io/gubernator/v2"
 	"go.opentelemetry.io/collector/component"
@@ -165,7 +167,7 @@ func (r *gubernatorRateLimiter) RateLimit(ctx context.Context, hits int) error {
 				zap.String("processor_id", r.set.ID.String()),
 				zap.Strings("metadata_keys", r.cfg.MetadataKeys),
 			)
-			return errTooManyRequests
+			return status.Error(codes.ResourceExhausted, errTooManyRequests.Error())
 		case ThrottleBehaviorDelay:
 			delay := time.Duration(resp.GetResetTime()-createdAt) * time.Millisecond
 			timer := time.NewTimer(delay)

processor/ratelimitprocessor/gubernator_test.go

@@ -100,7 +100,7 @@ func TestGubernatorRateLimiter_RateLimit(t *testing.T) {
 			err = rateLimiter.RateLimit(context.Background(), 1)
 			switch behavior {
 			case ThrottleBehaviorError:
-				assert.EqualError(t, err, "too many requests")
+				assert.EqualError(t, err, "rpc error: code = ResourceExhausted desc = too many requests")
 			case ThrottleBehaviorDelay:
 				assert.NoError(t, err)
 			}

processor/ratelimitprocessor/local.go

@@ -25,6 +25,8 @@ import (
 	"go.opentelemetry.io/collector/component"
 	"go.opentelemetry.io/collector/processor"
 	"golang.org/x/time/rate"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
 var _ RateLimiter = (*localRateLimiter)(nil)
@@ -61,12 +63,12 @@ func (r *localRateLimiter) RateLimit(ctx context.Context, hits int) error {
 	switch cfg.ThrottleBehavior {
 	case ThrottleBehaviorError:
 		if ok := limiter.AllowN(time.Now(), hits); !ok {
-			return errTooManyRequests
+			return status.Error(codes.ResourceExhausted, errTooManyRequests.Error())
 		}
 	case ThrottleBehaviorDelay:
 		lr := limiter.ReserveN(time.Now(), hits)
 		if !lr.OK() {
-			return errTooManyRequests
+			return status.Error(codes.ResourceExhausted, errTooManyRequests.Error())
 		}
 		timer := time.NewTimer(lr.Delay())
 		defer timer.Stop()

processor/ratelimitprocessor/local_test.go

@@ -80,7 +80,7 @@ func TestLocalRateLimiter_RateLimit(t *testing.T) {
 			err = rateLimiter.RateLimit(context.Background(), 1) // should fail
 			switch behavior {
 			case ThrottleBehaviorError:
-				assert.EqualError(t, err, "too many requests")
+				assert.EqualError(t, err, "rpc error: code = ResourceExhausted desc = too many requests")
 				// retry every 20ms to ensure that RateLimit will recover from error when bucket refills after 1 second
 				assert.Eventually(t, func() bool {
 					return rateLimiter.RateLimit(context.Background(), 1) == nil

processor/ratelimitprocessor/processor.go

@@ -19,7 +19,6 @@ package ratelimitprocessor // import "github.com/elastic/opentelemetry-collector
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"sync/atomic"
 	"time"
@@ -32,6 +31,8 @@ import (
 	"go.opentelemetry.io/collector/pdata/ptrace"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/metric"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 
 	"github.com/elastic/opentelemetry-collector-components/internal/sharedcomponent"
 	"github.com/elastic/opentelemetry-collector-components/processor/ratelimitprocessor/internal/metadata"
@@ -197,7 +198,7 @@ func getTelemetryAttrs(attrsCommon []attribute.KeyValue, err error) (attrs []att
 			telemetry.WithReason(telemetry.StatusUnderLimit),
 			telemetry.WithDecision("accepted"),
 		)
-	case errors.Is(err, errTooManyRequests):
+	case status.Code(err) == codes.ResourceExhausted:
 		attrs = append(attrsCommon,
 			telemetry.WithDecision("throttled"),
 		)
@@ -211,8 +212,7 @@ func getTelemetryAttrs(attrsCommon []attribute.KeyValue, err error) (attrs []att
 	return attrs
 }
 
-func rateLimit(
-	ctx context.Context,
+func rateLimit(ctx context.Context,
 	hits int,
 	rateLimit func(ctx context.Context, n int) error,
 	metadataKeys []string,

processor/ratelimitprocessor/processor_test.go

@@ -164,7 +164,7 @@ func TestConsume_Logs(t *testing.T) {
 	consumed = false
 	err = processor.ConsumeLogs(clientContext, logs)
 	assert.False(t, consumed)
-	assert.EqualError(t, err, "too many requests")
+	assert.EqualError(t, err, "rpc error: code = ResourceExhausted desc = too many requests")
 
 	testRateLimitTelemetry(t, tt)
 }
@@ -211,7 +211,7 @@ func TestConsume_Metrics(t *testing.T) {
 	consumed = false
 	err = processor.ConsumeMetrics(clientContext, metrics)
 	assert.False(t, consumed)
-	assert.EqualError(t, err, "too many requests")
+	assert.EqualError(t, err, "rpc error: code = ResourceExhausted desc = too many requests")
 
 	testRateLimitTelemetry(t, tt)
 }
@@ -258,7 +258,7 @@ func TestConsume_Traces(t *testing.T) {
 	consumed = false
 	err = processor.ConsumeTraces(clientContext, traces)
 	assert.False(t, consumed)
-	assert.EqualError(t, err, "too many requests")
+	assert.EqualError(t, err, "rpc error: code = ResourceExhausted desc = too many requests")
 
 	testRateLimitTelemetry(t, tt)
 }
@@ -306,7 +306,7 @@ func TestConsume_Profiles(t *testing.T) {
 	consumed = false
 	err = processor.ConsumeProfiles(clientContext, profiles)
 	assert.False(t, consumed)
-	assert.EqualError(t, err, "too many requests")
+	assert.EqualError(t, err, "rpc error: code = ResourceExhausted desc = too many requests")
 
 	testRateLimitTelemetry(t, tt)
 }