diff --git a/rest/handler/timeouthandler.go b/rest/handler/timeouthandler.go index 52d82a80..38a68122 100644 --- a/rest/handler/timeouthandler.go +++ b/rest/handler/timeouthandler.go @@ -128,7 +128,13 @@ type timeoutWriter struct { var _ http.Pusher = (*timeoutWriter)(nil) func (tw *timeoutWriter) Flush() { + dst := tw.w.Header() + for k, vv := range tw.h { + dst[k] = vv + } if flusher, ok := tw.w.(http.Flusher); ok { + tw.w.Write(tw.wbuf.Bytes()) + tw.wbuf.Reset() flusher.Flush() } } diff --git a/rest/handler/timeouthandler_test.go b/rest/handler/timeouthandler_test.go index acbda030..4f496fe0 100644 --- a/rest/handler/timeouthandler_test.go +++ b/rest/handler/timeouthandler_test.go @@ -1,9 +1,13 @@ package handler import ( + "bufio" "context" + "fmt" "net/http" "net/http/httptest" + "strconv" + "strings" "testing" "time" @@ -12,6 +16,38 @@ import ( "github.com/zeromicro/go-zero/rest/internal/response" ) +func TestTimeoutWriteFlushOutput(t *testing.T) { + timeoutHandler := TimeoutHandler(1000 * time.Millisecond) + handler := timeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream;charset=utf-8") + flusher, ok := w.(http.Flusher) + if !ok { + http.Error(w, "Flushing not supported", http.StatusInternalServerError) + return + } + for i := 1; i <= 5; i++ { + fmt.Fprint(w, strconv.Itoa(i)+"只猫猫\n\n") + flusher.Flush() + time.Sleep(time.Millisecond) + } + })) + req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody) + resp := httptest.NewRecorder() + handler.ServeHTTP(resp, req) + scanner := bufio.NewScanner(resp.Body) + mao := 0 + for scanner.Scan() { + line := scanner.Text() + if strings.Contains(line, "猫猫") { + mao++ + } + } + if err := scanner.Err(); err != nil { + mao = 0 + } + assert.Equal(t, "5只猫猫", strconv.Itoa(mao)+"只猫猫") +} + func TestTimeout(t *testing.T) { timeoutHandler := TimeoutHandler(time.Millisecond) handler := timeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {