fix: do not warn if requestid middleware errors due to ErrIllegalHeaderWrite (#2654)

Author: miparnisari

internal/middleware/usagemetrics/usagemetrics.go

@@ -3,6 +3,7 @@ package usagemetrics
 import (
 	"context"
 	"strconv"
+	"strings"
 	"time"
 
 	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors"
@@ -38,7 +39,7 @@ type reporter struct{}
 
 func (r *reporter) ServerReporter(ctx context.Context, callMeta interceptors.CallMeta) (interceptors.Reporter, context.Context) {
 	_, methodName := grpcutil.SplitMethodName(callMeta.FullMethod())
-	ctx = ContextWithHandle(ctx)
+	ctx = contextWithHandle(ctx)
 	return &serverReporter{ctx: ctx, methodName: methodName}, ctx
 }
 
@@ -48,19 +49,25 @@ type serverReporter struct {
 	methodName string
 }
 
+// PostCall is invoked after all PostMsgSend operations.
 func (r *serverReporter) PostCall(_ error, _ time.Duration) {
 	responseMeta := FromContext(r.ctx)
 	if responseMeta == nil {
 		responseMeta = &dispatch.ResponseMeta{}
 	}
 
-	err := annotateAndReportForMetadata(r.ctx, r.methodName, responseMeta)
-	// if context is cancelled, the stream will be closed, and gRPC will return ErrIllegalHeaderWrite
-	// this prevents logging unnecessary error messages
-	if r.ctx.Err() != nil {
-		return
-	}
+	DispatchedCountHistogram.WithLabelValues(r.methodName, "false").Observe(float64(responseMeta.DispatchCount))
+	DispatchedCountHistogram.WithLabelValues(r.methodName, "true").Observe(float64(responseMeta.CachedDispatchCount))
+	err := responsemeta.SetResponseTrailerMetadata(r.ctx, map[responsemeta.ResponseMetadataTrailerKey]string{
+		responsemeta.DispatchedOperationsCount: strconv.Itoa(int(responseMeta.DispatchCount)),
+		responsemeta.CachedOperationsCount:     strconv.Itoa(int(responseMeta.CachedDispatchCount)),
+	})
 	if err != nil {
+		// if context is cancelled, the stream will be closed, and gRPC will return ErrIllegalHeaderWrite (which is private)
+		// this prevents logging unnecessary error messages
+		if strings.Contains(err.Error(), "SendHeader called multiple times") {
+			return
+		}
 		log.Ctx(r.ctx).Warn().Err(err).Msg("usagemetrics: could not report metadata")
 	}
 }
@@ -79,16 +86,6 @@ func StreamServerInterceptor() grpc.StreamServerInterceptor {
 	return interceptors.StreamServerInterceptor(&reporter{})
 }
 
-func annotateAndReportForMetadata(ctx context.Context, methodName string, metadata *dispatch.ResponseMeta) error {
-	DispatchedCountHistogram.WithLabelValues(methodName, "false").Observe(float64(metadata.DispatchCount))
-	DispatchedCountHistogram.WithLabelValues(methodName, "true").Observe(float64(metadata.CachedDispatchCount))
-
-	return responsemeta.SetResponseTrailerMetadata(ctx, map[responsemeta.ResponseMetadataTrailerKey]string{
-		responsemeta.DispatchedOperationsCount: strconv.Itoa(int(metadata.DispatchCount)),
-		responsemeta.CachedOperationsCount:     strconv.Itoa(int(metadata.CachedDispatchCount)),
-	})
-}
-
 // Create a new type to prevent context collisions
 type responseMetaKey string
 
@@ -119,11 +116,11 @@ func FromContext(ctx context.Context) *dispatch.ResponseMeta {
 	return possibleHandle.(*metaHandle).metadata
 }
 
-// ContextWithHandle creates a new context with a location to store metadata
+// contextWithHandle creates a new context with a location to store metadata
 // returned from a dispatched request.
 //
 // This should only be called in middleware or testing functions.
-func ContextWithHandle(ctx context.Context) context.Context {
+func contextWithHandle(ctx context.Context) context.Context {
 	var handle metaHandle
 	return context.WithValue(ctx, metadataCtxKey, &handle)
 }

internal/middleware/usagemetrics/usagemetrics_test.go

@@ -3,7 +3,6 @@ package usagemetrics
 import (
 	"context"
 	"errors"
-	"fmt"
 	"io"
 	"testing"
 
@@ -21,53 +20,28 @@ type testServer struct {
 	testpb.UnimplementedTestServiceServer
 }
 
-func (t testServer) PingEmpty(ctx context.Context, _ *testpb.PingEmptyRequest) (*testpb.PingEmptyResponse, error) {
-	SetInContext(ctx, &dispatch.ResponseMeta{
-		DispatchCount:       1,
-		CachedDispatchCount: 1,
-	})
-	return &testpb.PingEmptyResponse{}, nil
-}
-
-func (t testServer) Ping(ctx context.Context, _ *testpb.PingRequest) (*testpb.PingResponse, error) {
+func (t *testServer) Ping(ctx context.Context, _ *testpb.PingRequest) (*testpb.PingResponse, error) {
 	SetInContext(ctx, &dispatch.ResponseMeta{
 		DispatchCount:       1,
 		CachedDispatchCount: 1,
 	})
 	return &testpb.PingResponse{Value: ""}, nil
 }
 
-func (t testServer) PingError(ctx context.Context, _ *testpb.PingErrorRequest) (*testpb.PingErrorResponse, error) {
-	SetInContext(ctx, &dispatch.ResponseMeta{
-		DispatchCount:       1,
-		CachedDispatchCount: 1,
-	})
-	return nil, fmt.Errorf("err")
+// PingError returns the context error
+func (t *testServer) PingError(ctx context.Context, _ *testpb.PingErrorRequest) (*testpb.PingErrorResponse, error) {
+	<-ctx.Done()
+	return nil, ctx.Err()
 }
 
-func (t testServer) PingList(_ *testpb.PingListRequest, server testpb.TestService_PingListServer) error {
+func (t *testServer) PingList(_ *testpb.PingListRequest, server testpb.TestService_PingListServer) error {
 	SetInContext(server.Context(), &dispatch.ResponseMeta{
 		DispatchCount:       1,
 		CachedDispatchCount: 1,
 	})
 	return nil
 }
 
-func (t testServer) PingStream(stream testpb.TestService_PingStreamServer) error {
-	count := int32(0)
-	for {
-		_, err := stream.Recv()
-		if errors.Is(err, io.EOF) {
-			break
-		} else if err != nil {
-			return err
-		}
-		_ = stream.Send(&testpb.PingStreamResponse{Value: "", Counter: count})
-		count++
-	}
-	return nil
-}
-
 type metricsMiddlewareTestSuite struct {
 	*testpb.InterceptorTestSuite
 }
@@ -131,3 +105,29 @@ func (s *metricsMiddlewareTestSuite) TestTrailers_Stream() {
 	s.Require().NoError(err)
 	s.Require().Equal(1, cachedCount)
 }
+
+func (s *metricsMiddlewareTestSuite) TestErrCtx() {
+	var trailerMD metadata.MD
+
+	// SimpleCtx times out after two seconds
+	_, err := s.Client.PingError(s.SimpleCtx(), &testpb.PingErrorRequest{}, grpc.Trailer(&trailerMD))
+	s.Require().ErrorContains(err, context.DeadlineExceeded.Error())
+
+	// TODO ideally, this test would assert that no error log has been written
+	// but right now we have no way of capturing the logs
+
+	// No metadata should have been sent
+	dispatchCount, err := responsemeta.GetIntResponseTrailerMetadata(
+		trailerMD,
+		responsemeta.DispatchedOperationsCount,
+	)
+	s.Require().ErrorContains(err, "key `io.spicedb.respmeta.dispatchedoperationscount` not found in trailer")
+	s.Require().Equal(0, dispatchCount)
+
+	cachedCount, err := responsemeta.GetIntResponseTrailerMetadata(
+		trailerMD,
+		responsemeta.CachedOperationsCount,
+	)
+	s.Require().ErrorContains(err, "key `io.spicedb.respmeta.cachedoperationscount` not found in trailer")
+	s.Require().Equal(0, cachedCount)
+}

pkg/middleware/requestid/requestid.go

@@ -2,6 +2,7 @@ package requestid
 
 import (
 	"context"
+	"strings"
 
 	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors"
 	"github.com/rs/xid"
@@ -54,6 +55,7 @@ func (r *handleRequestID) ClientReporter(ctx context.Context, meta interceptors.
 	return interceptors.NoopReporter{}, ctx
 }
 
+// ServerReporter is invoked before the request begins processing.
 func (r *handleRequestID) ServerReporter(ctx context.Context, _ interceptors.CallMeta) (interceptors.Reporter, context.Context) {
 	haveRequestID, requestID, ctx := r.fromContextOrGenerate(ctx)
 
@@ -67,7 +69,11 @@ func (r *handleRequestID) ServerReporter(ctx context.Context, _ interceptors.Cal
 			responsemeta.ResponseMetadataTrailerKey(responsemeta.RequestID): requestID,
 		})
 		if err != nil {
-			log.Ctx(ctx).Warn().Err(err).Msg("requestid: could not report metadata")
+			// if context is cancelled, the stream will be closed, and gRPC will return ErrIllegalHeaderWrite (which is private)
+			// this prevents logging unnecessary error messages
+			if !strings.Contains(err.Error(), "SendHeader called multiple times") {
+				log.Ctx(ctx).Warn().Err(err).Msg("requestid: could not report metadata")
+			}
 		}
 	}
 

pkg/middleware/requestid/requestid_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"io"
+	"sync"
 	"testing"
 	"time"
 
@@ -358,17 +359,29 @@ func TestRequestIDMiddlewareTimeout(t *testing.T) {
 	suite.Run(t, s)
 }
 
+// If a stream receives a context that is already cancelled,
+// the middleware will not be invoked at all. Therefore, we try to induce
+// what happens when the context errors *while the middleware is running*.
+// However, right now, this test is only useful when ran by hand and inspecting logs manually,
+// because we don't have a way of asserting what logs were emmitted by the middleware.
 func (s *requestIDTimeoutTestSuite) TestContextTimeout() {
-	// Create a context that's already cancelled
-	ctx, cancel := context.WithTimeout(s.SimpleCtx(), time.Nanosecond)
-	defer cancel()
-
-	// Wait for context to be cancelled
-	time.Sleep(time.Millisecond)
-
-	var trailer metadata.MD
-	_, err := s.Client.PingEmpty(ctx, &testpb.PingEmptyRequest{}, grpc.Trailer(&trailer))
+	var wg sync.WaitGroup
+	for i := 0; i < 10_000; i++ {
+		wg.Add(1)
+		wg.Go(func() {
+			defer wg.Done()
+
+			// context will be cancelled in the middle of the middleware (if we are lucky)
+			ctx, cancel := context.WithTimeout(s.T().Context(), 1*time.Millisecond)
+			defer cancel()
+
+			var trailer metadata.MD
+			_, err := s.Client.PingEmpty(ctx, &testpb.PingEmptyRequest{}, grpc.Trailer(&trailer))
+
+			// The RPC should fail due to context cancellation, but middleware should handle it gracefully
+			require.Error(s.T(), err)
+		})
+	}
 
-	// The RPC should fail due to context cancellation, but middleware should handle it gracefully
-	require.Error(s.T(), err)
+	wg.Wait()
 }

pkg/middleware/serverversion/serverversion.go

@@ -37,7 +37,7 @@ func (r *HandleServerVersion) ServerReporter(ctx context.Context, _ interceptors
 				err = responsemeta.SetResponseTrailerMetadata(ctx, map[responsemeta.ResponseMetadataTrailerKey]string{
 					responsemeta.ResponseMetadataTrailerKey(responsemeta.ServerVersion): version,
 				})
-				// if context is cancelled, the stream will be closed, and gRPC will return ErrIllegalHeaderWrite
+				// if context is cancelled, the stream will be closed, and gRPC will return ErrIllegalHeaderWrite (which is private)
 				// this prevents logging unnecessary error messages
 				if err := ctx.Err(); err != nil {
 					return interceptors.NoopReporter{}, ctx