Skip to content
This repository has been archived by the owner on Nov 5, 2024. It is now read-only.

Commit

Permalink
代理:如果是event-stream类型的请求自动Flush收到的数据
Browse files Browse the repository at this point in the history
  • Loading branch information
iwind committed Sep 8, 2020
1 parent 62a5c14 commit 229b71d
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 1 deletion.
19 changes: 18 additions & 1 deletion internal/teaproxy/request_backend.go
Original file line number Diff line number Diff line change
Expand Up @@ -234,6 +234,9 @@ func (this *Request) callBackend(writer *ResponseWriter) error {
this.responseCallback(writer)
}

// 是否需要刷新
shouldFlush := this.raw.Header.Get("Accept") == "text/event-stream"

// 准备
writer.Prepare(resp.ContentLength)

Expand All @@ -242,7 +245,21 @@ func (this *Request) callBackend(writer *ResponseWriter) error {

pool := this.bytePool(resp.ContentLength)
buf := pool.Get()
_, err = io.CopyBuffer(writer, resp.Body, buf)
if shouldFlush {
for {
n, readErr := resp.Body.Read(buf)
if n > 0 {
_, err = writer.Write(buf[:n])
writer.Flush()
}
if readErr != nil {
err = readErr
break
}
}
} else {
_, err = io.CopyBuffer(writer, resp.Body, buf)
}
pool.Put(buf)

err1 := resp.Body.Close()
Expand Down
8 changes: 8 additions & 0 deletions internal/teaproxy/response.go
Original file line number Diff line number Diff line change
Expand Up @@ -236,3 +236,11 @@ func (this *ResponseWriter) Hijack() (conn net.Conn, buf *bufio.ReadWriter, err
}
return
}

// Flush
func (this *ResponseWriter) Flush() {
flusher, ok := this.writer.(http.Flusher)
if ok {
flusher.Flush()
}
}

0 comments on commit 229b71d

Please sign in to comment.