integration: fix wrapping of response writer causing endless test loop

docs/docs/operations/logging.md

@@ -22,6 +22,12 @@ func (rsp *logResponseWriter) WriteHeader(code int) {
 	rsp.ResponseWriter.WriteHeader(code)
 }
 
+// Unwrap returns the original http.ResponseWriter. This is necessary
+// to expose Flush() and Push() on the underlying response writer.
+func (rsp *logResponseWriter) Unwrap() http.ResponseWriter {
+	return rsp.ResponseWriter
+}
+
 func newLogResponseWriter(w http.ResponseWriter) *logResponseWriter {
 	return &logResponseWriter{w, http.StatusOK}
 }
@@ -58,4 +64,4 @@ func logRequestBody(h http.Handler) http.Handler {
     s := &http.Server{
         Handler: logRequestBody(mux),
     }
-```
+```

examples/internal/gateway/handlers.go

@@ -77,6 +77,12 @@ func (rsp *logResponseWriter) WriteHeader(code int) {
 	rsp.ResponseWriter.WriteHeader(code)
 }
 
+// Unwrap returns the original http.ResponseWriter. This is necessary
+// to expose Flush() and Push() on the underlying response writer.
+func (rsp *logResponseWriter) Unwrap() http.ResponseWriter {
+	return rsp.ResponseWriter
+}
+
 func newLogResponseWriter(w http.ResponseWriter) *logResponseWriter {
 	return &logResponseWriter{w, http.StatusOK}
 }

runtime/handler.go

@@ -17,13 +17,8 @@ import (
 
 // ForwardResponseStream forwards the stream from gRPC server to REST client.
 func ForwardResponseStream(ctx context.Context, mux *ServeMux, marshaler Marshaler, w http.ResponseWriter, req *http.Request, recv func() (proto.Message, error), opts ...func(context.Context, http.ResponseWriter, proto.Message) error) {
-	f, ok := w.(http.Flusher)
-	if !ok {
-		grpclog.Infof("Flush not supported in %T", w)
-		http.Error(w, "unexpected type of web server", http.StatusInternalServerError)
-		return
-	}
-
+	// nolint: bodyclose
+	rc := http.NewResponseController(w)
 	md, ok := ServerMetadataFromContext(ctx)
 	if !ok {
 		grpclog.Infof("Failed to extract ServerMetadata from context")
@@ -94,7 +89,16 @@ func ForwardResponseStream(ctx context.Context, mux *ServeMux, marshaler Marshal
 			grpclog.Infof("Failed to send delimiter chunk: %v", err)
 			return
 		}
-		f.Flush()
+		err = rc.Flush()
+		if err != nil {
+			if errors.Is(err, http.ErrNotSupported) {
+				grpclog.Infof("Flush not supported in %T", w)
+				http.Error(w, "unexpected type of web server", http.StatusInternalServerError)
+				return
+			}
+			grpclog.Infof("Failed to flush response to client: %v", err)
+			return
+		}
 	}
 }