Introduce respWriter.Flush so we don't write the status twice

CHANGELOG.md

@@ -8,6 +8,10 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ## [Unreleased]
 
+### Fixed
+
+- The superfluous `response.WriteHeader` call in `go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp` when the response writer is flushed. (#5634)
+
 ## [1.27.0/0.52.0/0.21.0/0.7.0/0.2.0] - 2024-05-21
 
 ### Added

instrumentation/net/http/otelhttp/handler.go

@@ -204,6 +204,9 @@ func (h *middleware) serveHTTP(w http.ResponseWriter, r *http.Request, next http
 		WriteHeader: func(httpsnoop.WriteHeaderFunc) httpsnoop.WriteHeaderFunc {
 			return rww.WriteHeader
 		},
+		Flush: func(httpsnoop.FlushFunc) httpsnoop.FlushFunc {
+			return rww.Flush
+		},
 	})
 
 	labeler := &Labeler{}

instrumentation/net/http/otelhttp/handler_test.go

@@ -29,6 +29,9 @@ func TestHandler(t *testing.T) {
 				return otelhttp.NewHandler(
 					http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 						assert.Implements(t, (*http.Flusher)(nil), w)
+
+						w.(http.Flusher).Flush()
+						_, _ = io.WriteString(w, "Hello, world!\n")
 					}), "test_handler",
 				)
 			},

instrumentation/net/http/otelhttp/wrap.go

@@ -87,3 +87,13 @@ func (w *respWriterWrapper) WriteHeader(statusCode int) {
 	}
 	w.ResponseWriter.WriteHeader(statusCode)
 }
+
+func (w *respWriterWrapper) Flush() {
+	if !w.wroteHeader {
+		w.WriteHeader(http.StatusOK)
+	}
+
+	if f, ok := w.ResponseWriter.(http.Flusher); ok {
+		f.Flush()
+	}
+}

instrumentation/net/http/otelhttp/wrap_test.go

@@ -0,0 +1,60 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package otelhttp
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRespWriterWriteHeader(t *testing.T) {
+	rw := &respWriterWrapper{
+		ResponseWriter: &httptest.ResponseRecorder{},
+		record:         func(int64) {},
+	}
+
+	rw.WriteHeader(http.StatusTeapot)
+	assert.Equal(t, http.StatusTeapot, rw.statusCode)
+	assert.True(t, rw.wroteHeader)
+
+	rw.WriteHeader(http.StatusGone)
+	assert.Equal(t, http.StatusTeapot, rw.statusCode)
+}
+
+func TestRespWriterFlush(t *testing.T) {
+	rw := &respWriterWrapper{
+		ResponseWriter: &httptest.ResponseRecorder{},
+		record:         func(int64) {},
+	}
+
+	rw.Flush()
+	assert.Equal(t, http.StatusOK, rw.statusCode)
+	assert.True(t, rw.wroteHeader)
+}
+
+type nonFlushableResponseWriter struct{}
+
+func (_ nonFlushableResponseWriter) Header() http.Header {
+	return http.Header{}
+}
+
+func (_ nonFlushableResponseWriter) Write([]byte) (int, error) {
+	return 0, nil
+}
+
+func (_ nonFlushableResponseWriter) WriteHeader(int) {}
+
+func TestRespWriterFlushNoFlusher(t *testing.T) {
+	rw := &respWriterWrapper{
+		ResponseWriter: nonFlushableResponseWriter{},
+		record:         func(int64) {},
+	}
+
+	rw.Flush()
+	assert.Equal(t, http.StatusOK, rw.statusCode)
+	assert.True(t, rw.wroteHeader)
+}