feat: support http stream response (#4055)

This commit is contained in:
Kevin Wan
2024-04-09 20:46:44 +08:00
committed by GitHub
parent 5fbe8ff5c4
commit 3866b5741a
2 changed files with 78 additions and 0 deletions

View File

@@ -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 {