From 94799103a321c114e7202f63bf95dcb940e0133a Mon Sep 17 00:00:00 2001
From: Damien Mathieu <42@dmathieu.com>
Date: Tue, 28 May 2024 20:10:49 +0200
Subject: [PATCH] Introduce respWriter.Flush so we don't write the status twice
 (#5634)

Closes #5438.
---
 CHANGELOG.md                                  |  4 ++
 instrumentation/net/http/otelhttp/handler.go  |  3 +
 .../net/http/otelhttp/handler_test.go         |  3 +
 instrumentation/net/http/otelhttp/wrap.go     | 10 ++++
 .../net/http/otelhttp/wrap_test.go            | 60 +++++++++++++++++++
 5 files changed, 80 insertions(+)
 create mode 100644 instrumentation/net/http/otelhttp/wrap_test.go

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0d2d722288d..36eff57b763 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,6 +12,10 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 - The `go.opentelemetry.io/contrib/config` add support to configure periodic reader interval and timeout. (#5661)
 
+### Fixed
+
+- The superfluous `response.WriteHeader` call in `go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp` when the response writer is flushed. (#5634)
+
 ### Deprecated
 
 - The `go.opentelemetry.io/contrib/instrumentation/go.mongodb.org/mongo-driver/mongo/otelmongo` package is deprecated.
diff --git a/instrumentation/net/http/otelhttp/handler.go b/instrumentation/net/http/otelhttp/handler.go
index c64f8beca71..810b7dd769d 100644
--- a/instrumentation/net/http/otelhttp/handler.go
+++ b/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{}
diff --git a/instrumentation/net/http/otelhttp/handler_test.go b/instrumentation/net/http/otelhttp/handler_test.go
index 0ba6eeeb985..8b0b7a23a3b 100644
--- a/instrumentation/net/http/otelhttp/handler_test.go
+++ b/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",
 				)
 			},
diff --git a/instrumentation/net/http/otelhttp/wrap.go b/instrumentation/net/http/otelhttp/wrap.go
index 2f4cc124dc6..948f8406c09 100644
--- a/instrumentation/net/http/otelhttp/wrap.go
+++ b/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()
+	}
+}
diff --git a/instrumentation/net/http/otelhttp/wrap_test.go b/instrumentation/net/http/otelhttp/wrap_test.go
new file mode 100644
index 00000000000..d4b89411a29
--- /dev/null
+++ b/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)
+}