Repository: zeromicro/go-zero
Pull request #4055 (2024-04-09 10:24:46)

feat: support http stream response
1. support `httpx.Stream` method
2. set `timeout` to 0 to use `httpx.Stream` if multiple writes cannot complete within `timeout`

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/zeromicro/go-zero/pull/4055?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zeromicro) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 95.16%. Comparing base [(`8690859`)](https://app.codecov.io/gh/zeromicro/go-zero/commit/8690859c7d1c9461bdc422c3c65608bde30904db?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zeromicro) to head [(`02c483c`)](https://app.codecov.io/gh/zeromicro/go-zero/pull/4055?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zeromicro).
> Report is 3 commits behind head on master.

<details><summary>Additional details and impacted files</summary>



| [Files](https://app.codecov.io/gh/zeromicro/go-zero/pull/4055?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zeromicro) | Coverage Δ | |
|---|---|---|
| [rest/httpx/responses.go](https://app.codecov.io/gh/zeromicro/go-zero/pull/4055?src=pr&el=tree&filepath=rest%2Fhttpx%2Fresponses.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zeromicro#diff-cmVzdC9odHRweC9yZXNwb25zZXMuZ28=) | `100.00% <100.00%> (ø)` | |

... and [254 files with indirect coverage changes](https://app.codecov.io/gh/zeromicro/go-zero/pull/4055/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zeromicro)

</details>

Patch:
diff --git a/rest/httpx/responses.go b/rest/httpx/responses.go
index edd133a4..9c3d4429 100644
--- a/rest/httpx/responses.go
+++ b/rest/httpx/responses.go
@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"net/http"
 	"sync"
 
@@ -88,6 +89,26 @@ func SetOkHandler(handler func(context.Context, any) any) {
 	okHandler = handler
 }
 
+// Stream writes data into w with streaming mode.
+// The ctx is used to control the streaming loop, typically use r.Context().
+// The fn is called repeatedly until it returns false.
+func Stream(ctx context.Context, w http.ResponseWriter, fn func(w io.Writer) bool) {
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		default:
+			hasMore := fn(w)
+			if flusher, ok := w.(http.Flusher); ok {
+				flusher.Flush()
+			}
+			if !hasMore {
+				return
+			}
+		}
+	}
+}
+
 // WriteJson writes v as json string into w with code.
 func WriteJson(w http.ResponseWriter, code int, v any) {
 	if err := doWriteJson(w, code, v); err != nil {


Test patch:
diff --git a/rest/httpx/responses_test.go b/rest/httpx/responses_test.go
index 44d4d720..00e099da 100644
--- a/rest/httpx/responses_test.go
+++ b/rest/httpx/responses_test.go
@@ -1,10 +1,13 @@
 package httpx
 
 import (
+	"bytes"
 	"context"
 	"errors"
 	"fmt"
+	"io"
 	"net/http"
+	"net/http/httptest"
 	"strings"
 	"testing"
 
@@ -239,6 +242,60 @@ func TestWriteJsonMarshalFailed(t *testing.T) {
 	assert.Equal(t, http.StatusInternalServerError, w.code)
 }
 
+func TestStream(t *testing.T) {
+	t.Run("regular case", func(t *testing.T) {
+		channel := make(chan string)
+		go func() {
+			defer close(channel)
+			for index := 0; index < 5; index++ {
+				channel <- fmt.Sprintf("%d", index)
+			}
+		}()
+
+		w := httptest.NewRecorder()
+		Stream(context.Background(), w, func(w io.Writer) bool {
+			output, ok := <-channel
+			if !ok {
+				return false
+			}
+
+			outputBytes := bytes.NewBufferString(output)
+			_, err := w.Write(append(outputBytes.Bytes(), []byte("\n")...))
+			return err == nil
+		})
+
+		assert.Equal(t, http.StatusOK, w.Code)
+		assert.Equal(t, "0\n1\n2\n3\n4\n", w.Body.String())
+	})
+
+	t.Run("context done", func(t *testing.T) {
+		channel := make(chan string)
+		go func() {
+			defer close(channel)
+			for index := 0; index < 5; index++ {
+				channel <- fmt.Sprintf("num: %d", index)
+			}
+		}()
+
+		w := httptest.NewRecorder()
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+		Stream(ctx, w, func(w io.Writer) bool {
+			output, ok := <-channel
+			if !ok {
+				return false
+			}
+
+			outputBytes := bytes.NewBufferString(output)
+			_, err := w.Write(append(outputBytes.Bytes(), []byte("\n")...))
+			return err == nil
+		})
+
+		assert.Equal(t, http.StatusOK, w.Code)
+		assert.Equal(t, "", w.Body.String())
+	})
+}
+
 type tracedResponseWriter struct {
 	headers     map[string][]string
 	builder     strings.Builder