Merge pull request #65 from tbphp/refactor-stream-response
refactor: 优化流式响应的性能和兼容性
This commit is contained in:
@@ -1,10 +1,8 @@
|
|||||||
package proxy
|
package proxy
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"bufio"
|
|
||||||
"net/http"
|
|
||||||
|
|
||||||
"io"
|
"io"
|
||||||
|
"net/http"
|
||||||
|
|
||||||
"github.com/gin-gonic/gin"
|
"github.com/gin-gonic/gin"
|
||||||
"github.com/sirupsen/logrus"
|
"github.com/sirupsen/logrus"
|
||||||
@@ -23,29 +21,23 @@ func (ps *ProxyServer) handleStreamingResponse(c *gin.Context, resp *http.Respon
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
scanner := bufio.NewScanner(resp.Body)
|
buf := make([]byte, 4*1024)
|
||||||
scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
|
for {
|
||||||
for scanner.Scan() {
|
n, err := resp.Body.Read(buf)
|
||||||
select {
|
if n > 0 {
|
||||||
case <-c.Request.Context().Done():
|
if _, writeErr := c.Writer.Write(buf[:n]); writeErr != nil {
|
||||||
logrus.Debugf("Client disconnected, closing stream.")
|
logUpstreamError("writing stream to client", writeErr)
|
||||||
return
|
return
|
||||||
default:
|
}
|
||||||
|
flusher.Flush()
|
||||||
}
|
}
|
||||||
|
if err == io.EOF {
|
||||||
if _, err := c.Writer.Write(scanner.Bytes()); err != nil {
|
break
|
||||||
logUpstreamError("writing stream to client", err)
|
}
|
||||||
|
if err != nil {
|
||||||
|
logUpstreamError("reading from upstream", err)
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
if _, err := c.Writer.Write([]byte("\n")); err != nil {
|
|
||||||
logUpstreamError("writing stream newline to client", err)
|
|
||||||
return
|
|
||||||
}
|
|
||||||
flusher.Flush()
|
|
||||||
}
|
|
||||||
|
|
||||||
if err := scanner.Err(); err != nil {
|
|
||||||
logUpstreamError("reading from upstream scanner", err)
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
Reference in New Issue
Block a user