Move otelhttp wrappers into internal package (#5916)

This moves the otelhttp wrappers into the internal package, so we cannot
rely on private values, and are forced to think about race conditions
when using that data.

Closes #5914.

---------

Co-authored-by: Tyler Yahn <[email protected]>

CHANGELOG.md

@@ -24,6 +24,10 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 - The deprecated `go.opentelemetry.io/contrib/processors/baggagecopy` package is removed. (#5853)
 
+### Fixed
+
+- Race condition when reading the HTTP body and writing the response in `go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp`. (#5916)
+
 <!-- Released section -->
 <!-- Don't change this section unless doing release -->
 

instrumentation/net/http/otelhttp/handler.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/felixge/httpsnoop"
 
+	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp/internal/request"
 	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp/internal/semconv"
 	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp/internal/semconvutil"
 	"go.opentelemetry.io/otel"
@@ -166,14 +167,12 @@ func (h *middleware) serveHTTP(w http.ResponseWriter, r *http.Request, next http
 		}
 	}
 
-	var bw bodyWrapper
 	// if request body is nil or NoBody, we don't want to mutate the body as it
 	// will affect the identity of it in an unforeseeable way because we assert
 	// ReadCloser fulfills a certain interface and it is indeed nil or NoBody.
+	bw := request.NewBodyWrapper(r.Body, readRecordFunc)
 	if r.Body != nil && r.Body != http.NoBody {
-		bw.ReadCloser = r.Body
-		bw.record = readRecordFunc
-		r.Body = &bw
+		r.Body = bw
 	}
 
 	writeRecordFunc := func(int64) {}
@@ -183,13 +182,7 @@ func (h *middleware) serveHTTP(w http.ResponseWriter, r *http.Request, next http
 		}
 	}
 
-	rww := &respWriterWrapper{
-		ResponseWriter: w,
-		record:         writeRecordFunc,
-		ctx:            ctx,
-		props:          h.propagators,
-		statusCode:     http.StatusOK, // default status code in case the Handler doesn't write anything
-	}
+	rww := request.NewRespWriterWrapper(w, writeRecordFunc)
 
 	// Wrap w to use our ResponseWriter methods while also exposing
 	// other interfaces that w may implement (http.CloseNotifier,
@@ -217,24 +210,26 @@ func (h *middleware) serveHTTP(w http.ResponseWriter, r *http.Request, next http
 
 	next.ServeHTTP(w, r.WithContext(ctx))
 
-	span.SetStatus(semconv.ServerStatus(rww.statusCode))
+	statusCode := rww.StatusCode()
+	bytesWritten := rww.BytesWritten()
+	span.SetStatus(semconv.ServerStatus(statusCode))
 	span.SetAttributes(h.traceSemconv.ResponseTraceAttrs(semconv.ResponseTelemetry{
-		StatusCode: rww.statusCode,
-		ReadBytes:  bw.read.Load(),
-		ReadError:  bw.err,
-		WriteBytes: rww.written,
-		WriteError: rww.err,
+		StatusCode: statusCode,
+		ReadBytes:  bw.BytesRead(),
+		ReadError:  bw.Error(),
+		WriteBytes: bytesWritten,
+		WriteError: rww.Error(),
 	})...)
 
 	// Add metrics
 	attributes := append(labeler.Get(), semconvutil.HTTPServerRequestMetrics(h.server, r)...)
-	if rww.statusCode > 0 {
-		attributes = append(attributes, semconv.HTTPStatusCode(rww.statusCode))
+	if statusCode > 0 {
+		attributes = append(attributes, semconv.HTTPStatusCode(statusCode))
 	}
 	o := metric.WithAttributeSet(attribute.NewSet(attributes...))
 
-	h.requestBytesCounter.Add(ctx, bw.read.Load(), o)
-	h.responseBytesCounter.Add(ctx, rww.written, o)
+	h.requestBytesCounter.Add(ctx, bw.BytesRead(), o)
+	h.responseBytesCounter.Add(ctx, bytesWritten, o)
 
 	// Use floating point division here for higher precision (instead of Millisecond method).
 	elapsedTime := float64(time.Since(requestStartTime)) / float64(time.Millisecond)

instrumentation/net/http/otelhttp/internal/request/body_wrapper.go

@@ -0,0 +1,75 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package request // import "go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp/internal/request"
+
+import (
+	"io"
+	"sync"
+)
+
+var _ io.ReadCloser = &BodyWrapper{}
+
+// BodyWrapper wraps a http.Request.Body (an io.ReadCloser) to track the number
+// of bytes read and the last error.
+type BodyWrapper struct {
+	io.ReadCloser
+	OnRead func(n int64) // must not be nil
+
+	mu   sync.Mutex
+	read int64
+	err  error
+}
+
+// NewBodyWrapper creates a new BodyWrapper.
+//
+// The onRead attribute is a callback that will be called every time the data
+// is read, with the number of bytes being read.
+func NewBodyWrapper(body io.ReadCloser, onRead func(int64)) *BodyWrapper {
+	return &BodyWrapper{
+		ReadCloser: body,
+		OnRead:     onRead,
+	}
+}
+
+// Read reads the data from the io.ReadCloser, and stores the number of bytes
+// read and the error.
+func (w *BodyWrapper) Read(b []byte) (int, error) {
+	n, err := w.ReadCloser.Read(b)
+	n1 := int64(n)
+
+	w.updateReadData(n1, err)
+	w.OnRead(n1)
+	return n, err
+}
+
+func (w *BodyWrapper) updateReadData(n int64, err error) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+
+	w.read += n
+	if err != nil {
+		w.err = err
+	}
+}
+
+// Closes closes the io.ReadCloser.
+func (w *BodyWrapper) Close() error {
+	return w.ReadCloser.Close()
+}
+
+// BytesRead returns the number of bytes read up to this point.
+func (w *BodyWrapper) BytesRead() int64 {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+
+	return w.read
+}
+
+// Error returns the last error.
+func (w *BodyWrapper) Error() error {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+
+	return w.err
+}

instrumentation/net/http/otelhttp/internal/request/body_wrapper_test.go

@@ -0,0 +1,74 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package request
+
+import (
+	"errors"
+	"io"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+var errFirstCall = errors.New("first call")
+
+func TestBodyWrapper(t *testing.T) {
+	bw := NewBodyWrapper(io.NopCloser(strings.NewReader("hello world")), func(int64) {})
+
+	data, err := io.ReadAll(bw)
+	require.NoError(t, err)
+	assert.Equal(t, "hello world", string(data))
+
+	assert.Equal(t, int64(11), bw.BytesRead())
+	assert.Equal(t, io.EOF, bw.Error())
+}
+
+type multipleErrorsReader struct {
+	calls int
+}
+
+type errorWrapper struct{}
+
+func (errorWrapper) Error() string {
+	return "subsequent calls"
+}
+
+func (mer *multipleErrorsReader) Read([]byte) (int, error) {
+	mer.calls = mer.calls + 1
+	if mer.calls == 1 {
+		return 0, errFirstCall
+	}
+
+	return 0, errorWrapper{}
+}
+
+func TestBodyWrapperWithErrors(t *testing.T) {
+	bw := NewBodyWrapper(io.NopCloser(&multipleErrorsReader{}), func(int64) {})
+
+	data, err := io.ReadAll(bw)
+	require.Equal(t, errFirstCall, err)
+	assert.Equal(t, "", string(data))
+	require.Equal(t, errFirstCall, bw.Error())
+
+	data, err = io.ReadAll(bw)
+	require.Equal(t, errorWrapper{}, err)
+	assert.Equal(t, "", string(data))
+	require.Equal(t, errorWrapper{}, bw.Error())
+}
+
+func TestConcurrentBodyWrapper(t *testing.T) {
+	bw := NewBodyWrapper(io.NopCloser(strings.NewReader("hello world")), func(int64) {})
+
+	go func() {
+		_, _ = io.ReadAll(bw)
+	}()
+
+	assert.NotNil(t, bw.BytesRead())
+	assert.Eventually(t, func() bool {
+		return errors.Is(bw.Error(), io.EOF)
+	}, time.Second, 10*time.Millisecond)
+}

instrumentation/net/http/otelhttp/internal/request/resp_writer_wrapper.go

@@ -0,0 +1,112 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package request // import "go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp/internal/request"
+
+import (
+	"net/http"
+	"sync"
+)
+
+var _ http.ResponseWriter = &RespWriterWrapper{}
+
+// RespWriterWrapper wraps a http.ResponseWriter in order to track the number of
+// bytes written, the last error, and to catch the first written statusCode.
+// TODO: The wrapped http.ResponseWriter doesn't implement any of the optional
+// types (http.Hijacker, http.Pusher, http.CloseNotifier, etc)
+// that may be useful when using it in real life situations.
+type RespWriterWrapper struct {
+	http.ResponseWriter
+	OnWrite func(n int64) // must not be nil
+
+	mu          sync.RWMutex
+	written     int64
+	statusCode  int
+	err         error
+	wroteHeader bool
+}
+
+// NewRespWriterWrapper creates a new RespWriterWrapper.
+//
+// The onWrite attribute is a callback that will be called every time the data
+// is written, with the number of bytes that were written.
+func NewRespWriterWrapper(w http.ResponseWriter, onWrite func(int64)) *RespWriterWrapper {
+	return &RespWriterWrapper{
+		ResponseWriter: w,
+		OnWrite:        onWrite,
+		statusCode:     http.StatusOK, // default status code in case the Handler doesn't write anything
+	}
+}
+
+// Write writes the bytes array into the [ResponseWriter], and tracks the
+// number of bytes written and last error.
+func (w *RespWriterWrapper) Write(p []byte) (int, error) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+
+	w.writeHeader(http.StatusOK)
+
+	n, err := w.ResponseWriter.Write(p)
+	n1 := int64(n)
+	w.OnWrite(n1)
+	w.written += n1
+	w.err = err
+	return n, err
+}
+
+// WriteHeader persists initial statusCode for span attribution.
+// All calls to WriteHeader will be propagated to the underlying ResponseWriter
+// and will persist the statusCode from the first call.
+// Blocking consecutive calls to WriteHeader alters expected behavior and will
+// remove warning logs from net/http where developers will notice incorrect handler implementations.
+func (w *RespWriterWrapper) WriteHeader(statusCode int) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+
+	w.writeHeader(statusCode)
+}
+
+// writeHeader persists the status code for span attribution, and propagates
+// the call to the underlying ResponseWriter.
+// It does not acquire a lock, and therefore assumes that is being handled by a
+// parent method.
+func (w *RespWriterWrapper) writeHeader(statusCode int) {
+	if !w.wroteHeader {
+		w.wroteHeader = true
+		w.statusCode = statusCode
+	}
+	w.ResponseWriter.WriteHeader(statusCode)
+}
+
+// Flush implements [http.Flusher].
+func (w *RespWriterWrapper) Flush() {
+	w.WriteHeader(http.StatusOK)
+
+	if f, ok := w.ResponseWriter.(http.Flusher); ok {
+		f.Flush()
+	}
+}
+
+// BytesWritten returns the number of bytes written.
+func (w *RespWriterWrapper) BytesWritten() int64 {
+	w.mu.RLock()
+	defer w.mu.RUnlock()
+
+	return w.written
+}
+
+// BytesWritten returns the HTTP status code that was sent.
+func (w *RespWriterWrapper) StatusCode() int {
+	w.mu.RLock()
+	defer w.mu.RUnlock()
+
+	return w.statusCode
+}
+
+// Error returns the last error.
+func (w *RespWriterWrapper) Error() error {
+	w.mu.RLock()
+	defer w.mu.RUnlock()
+
+	return w.err
+}