From 9e644ef111a33bc3c3fbbc1fca0d14a6ea6fff4f Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@gmail.com>
Date: Sun, 28 Feb 2021 19:21:30 +0200
Subject: [PATCH] lib/httpserver: make sure the gzipResponseWriter.Write() is
 called on Flush() and Close() calls

This should fix the `http: superfluous response.WriteHeader call` issue

See https://github.com/VictoriaMetrics/VictoriaMetrics/issues/1078
---
 docs/CHANGELOG.md            |  1 +
 lib/httpserver/httpserver.go | 27 ++++++++++++++++++---------
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md
index 0e9b6f695..8612c0b17 100644
--- a/docs/CHANGELOG.md
+++ b/docs/CHANGELOG.md
@@ -26,6 +26,7 @@
 * BUGFIX: vmselect: do not cache partial query results on timeout when receiving data from `vmstorage` nodes. See https://github.com/VictoriaMetrics/VictoriaMetrics/pull/1085
 * BUGFIX: properly handle `stale NFS file handle` error.
 * BUGFIX: properly cache query results when `extra_label` query arg is used. Previously the cached results could clash for different `extra_label` values. See https://github.com/VictoriaMetrics/VictoriaMetrics/pull/1095
+* BUGFIX: fix `http: superfluous response.WriteHeader call` issue. See https://github.com/VictoriaMetrics/VictoriaMetrics/issues/1078
 
 
 # [v1.54.1](https://github.com/VictoriaMetrics/VictoriaMetrics/releases/tag/v1.54.1)
diff --git a/lib/httpserver/httpserver.go b/lib/httpserver/httpserver.go
index d34ccd7e4..d09251e16 100644
--- a/lib/httpserver/httpserver.go
+++ b/lib/httpserver/httpserver.go
@@ -298,9 +298,9 @@ func maybeGzipResponseWriter(w http.ResponseWriter, r *http.Request) http.Respon
 	zw := getGzipWriter(w)
 	bw := getBufioWriter(zw)
 	zrw := &gzipResponseWriter{
-		ResponseWriter: w,
-		zw:             zw,
-		bw:             bw,
+		rw: w,
+		zw: zw,
+		bw: bw,
 	}
 	return zrw
 }
@@ -346,7 +346,7 @@ func putGzipWriter(zw *gzip.Writer) {
 var gzipWriterPool sync.Pool
 
 type gzipResponseWriter struct {
-	http.ResponseWriter
+	rw         http.ResponseWriter
 	zw         *gzip.Writer
 	bw         *bufio.Writer
 	statusCode int
@@ -355,6 +355,12 @@ type gzipResponseWriter struct {
 	disableCompression bool
 }
 
+// Implements http.ResponseWriter.Header method.
+func (zrw *gzipResponseWriter) Header() http.Header {
+	return zrw.rw.Header()
+}
+
+// Implements http.ResponseWriter.Write method.
 func (zrw *gzipResponseWriter) Write(p []byte) (int, error) {
 	if !zrw.firstWriteDone {
 		h := zrw.Header()
@@ -377,11 +383,12 @@ func (zrw *gzipResponseWriter) Write(p []byte) (int, error) {
 		zrw.firstWriteDone = true
 	}
 	if zrw.disableCompression {
-		return zrw.ResponseWriter.Write(p)
+		return zrw.rw.Write(p)
 	}
 	return zrw.bw.Write(p)
 }
 
+// Implements http.ResponseWriter.WriteHeader method.
 func (zrw *gzipResponseWriter) WriteHeader(statusCode int) {
 	zrw.statusCode = statusCode
 }
@@ -390,11 +397,14 @@ func (zrw *gzipResponseWriter) writeHeader() {
 	if zrw.statusCode == 0 {
 		zrw.statusCode = http.StatusOK
 	}
-	zrw.ResponseWriter.WriteHeader(zrw.statusCode)
+	zrw.rw.WriteHeader(zrw.statusCode)
 }
 
 // Implements http.Flusher
 func (zrw *gzipResponseWriter) Flush() {
+	if !zrw.firstWriteDone {
+		zrw.Write(nil)
+	}
 	if !zrw.disableCompression {
 		if err := zrw.bw.Flush(); err != nil && !isTrivialNetworkError(err) {
 			logger.Warnf("gzipResponseWriter.Flush (buffer): %s", err)
@@ -403,15 +413,14 @@ func (zrw *gzipResponseWriter) Flush() {
 			logger.Warnf("gzipResponseWriter.Flush (gzip): %s", err)
 		}
 	}
-	if fw, ok := zrw.ResponseWriter.(http.Flusher); ok {
+	if fw, ok := zrw.rw.(http.Flusher); ok {
 		fw.Flush()
 	}
 }
 
 func (zrw *gzipResponseWriter) Close() error {
 	if !zrw.firstWriteDone {
-		zrw.writeHeader()
-		return nil
+		zrw.Write(nil)
 	}
 	zrw.Flush()
 	var err error