fix: add flush to support chunked for the metric middleware (#2772)

Signed-off-by: saltbo <saltbo@foxmail.com>

pkg/utils/metrics/http_metrics.go

@@ -67,6 +67,12 @@ func (rw *ResponseWriterWrapper) WriteHeader(statuscode int) {
 	rw.ResponseWriter.WriteHeader(statuscode)
 }
 
+func (rw *ResponseWriterWrapper) Flush() {
+	if f, ok := rw.ResponseWriter.(http.Flusher); ok {
+		f.Flush()
+	}
+}
+
 func HTTPMetricMiddleware(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if util.IsWebsocketRequest(r) {

pkg/utils/metrics/http_metrics_test.go

@@ -0,0 +1,98 @@
+package metrics
+
+import (
+	"bufio"
+	"context"
+	"fmt"
+	"io"
+	"log"
+	"net/http"
+	"net/http/httptest"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/gorilla/mux"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+var dataRow = []byte("I'm the data Row\n")
+
+func chunkedHandler(w http.ResponseWriter, r *http.Request) {
+	ctx, cancel := context.WithCancel(r.Context())
+	ticker := time.NewTicker(time.Second) // We may set it to 10 secs
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		for {
+			select {
+			case <-ticker.C:
+				_, _ = w.Write(dataRow)
+				if f, ok := w.(http.Flusher); ok {
+					f.Flush()
+				}
+			case <-ctx.Done():
+				return
+			}
+		}
+	}()
+
+	// Emulate some work
+	time.Sleep(5 * time.Second)
+
+	// Telling the loop that keeps the connection alive to end
+	cancel()
+
+	// Waiting until the loop ends
+	wg.Wait()
+}
+
+func TestChunked(t *testing.T) {
+	mr := mux.NewRouter()
+	mr.Use(HTTPMetricMiddleware)
+	mr.Handle("/", http.HandlerFunc(chunkedHandler))
+	s := httptest.NewServer(mr)
+	defer s.Close()
+
+	resp, err := http.Get(s.URL)
+	require.NoError(t, err)
+	assert.Contains(t, resp.TransferEncoding, "chunked")
+	defer resp.Body.Close()
+
+	r := bufio.NewReader(resp.Body)
+	for {
+		line, err := readChunkedResponseLine(r)
+		if err != nil {
+			if err == io.EOF {
+				return
+			}
+			log.Fatal(err.Error())
+		}
+		if len(line) == 0 {
+			log.Println("Alive!")
+			continue
+		}
+
+		fmt.Println(string(line)) // we got the final response
+		assert.Equal(t, dataRow, append(line, '\n'))
+	}
+}
+
+func readChunkedResponseLine(r *bufio.Reader) ([]byte, error) {
+	line, isPrefix, err := r.ReadLine()
+	if err != nil {
+		return nil, err
+	}
+
+	if isPrefix {
+		rest, err := readChunkedResponseLine(r)
+		if err != nil {
+			return nil, err
+		}
+		line = append(line, rest...)
+	}
+
+	return line, nil
+}