[WIP] Implement proper payload chunked encoding in HTTP api_v3

cmd/query/app/apiv3/http_gateway.go

@@ -4,12 +4,14 @@
 package apiv3
 
 import (
+	"compress/gzip"
 	"encoding/json"
 	"errors"
 	"fmt"
 	"net/http"
 	"net/url"
 	"strconv"
+	"strings"
 	"time"
 
 	"github.com/gogo/protobuf/jsonpb"
@@ -51,6 +53,12 @@
 	Tracer       trace.TracerProvider
 }
 
+type chunkedGzipWriter struct {
+	http.ResponseWriter
+	gzipWriter *gzip.Writer
+	flusher    http.Flusher
+}
+
 // RegisterRoutes registers HTTP endpoints for APIv3 into provided mux.
 // The called can create a subrouter if it needs to prepend a base path.
 func (h *HTTPGateway) RegisterRoutes(router *mux.Router) {
@@ -60,15 +68,58 @@
 	h.addRoute(router, h.getOperations, routeGetOperations).Methods(http.MethodGet)
 }
 
-// addRoute adds a new endpoint to the router with given path and handler function.
-// This code is mostly copied from ../http_handler.
+func (w *chunkedGzipWriter) Write(b []byte) (int, error) {
+	n, err := w.gzipWriter.Write(b)
+	if err != nil {
+		return n, err
+	}
+
+	if err := w.gzipWriter.Flush(); err != nil {
+		return n, err
+	}
+
+	w.flusher.Flush()
+	return n, nil
+}
+
+func gzipMiddleware(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if !strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
+			next.ServeHTTP(w, r)
+			return
+		}
+
+		flusher, ok := w.(http.Flusher)
+		if !ok {
+			http.Error(w, "Streaming unsupported!", http.StatusInternalServerError)
+			return
+		}
+
+		// Set chunked transfer encoding before any writes
+		w.Header().Set("Transfer-Encoding", "chunked")
+		w.Header().Set("Content-Encoding", "gzip")
+		w.Header().Set("Vary", "Accept-Encoding")
+
+		gz := gzip.NewWriter(w)
+		defer gz.Close()
+
+		writer := &chunkedGzipWriter{
+			ResponseWriter: w,
+			gzipWriter:     gz,
+			flusher:        flusher,
+		}
+
+		next.ServeHTTP(writer, r)
+	})
+}
+
 func (*HTTPGateway) addRoute(
 	router *mux.Router,
 	f func(http.ResponseWriter, *http.Request),
 	route string,
-	_ ...any, /* args */
 ) *mux.Route {
 	var handler http.Handler = http.HandlerFunc(f)
+	handler = gzipMiddleware(handler)
 	handler = otelhttp.WithRouteTag(route, handler)
 	handler = spanNameHandler(route, handler)
 	return router.HandleFunc(route, handler.ServeHTTP)

cmd/query/app/server_test.go

@@ -7,6 +7,7 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"io"
 	"net"
 	"net/http"
 	"testing"
@@ -915,6 +916,21 @@ func TestServerHTTP_TracesRequest(t *testing.T) {
 			resp, err := client.Do(req)
 			require.NoError(t, err)
 
+			var fullResponse bytes.Buffer
+
+			buf := make([]byte, 1024)
+			for {
+				n, err := resp.Body.Read(buf)
+
+				if n > 0 {
+					fullResponse.Write(buf[:n])
+				}
+
+				if err == io.EOF {
+					break
+				}
+			}
+
 			if test.expectedTrace != "" {
 				assert.Len(t, exporter.GetSpans(), 1, "HTTP request was traced and span reported")
 				assert.Equal(t, test.expectedTrace, exporter.GetSpans()[0].Name)