feat: support native streaming responses

2026-09-21 04:02:04 +08:00
parent 694374fde8
commit 62b0f09546
6 changed files with 252 additions and 9 deletions
--- a/transport/native/server.go
+++ b/transport/native/server.go
@@ -334,6 +334,7 @@ func (s *Server) serveConnection(baseCtx context.Context, state *connection) {
 		state.active.Store(true)
 		writer := s.writerPool.Get().(*responseWriter)
 		writer.reset(request.Method)
+		writer.configure(request, state.conn, closeAfter, s.WriteTimeout)
 		requestSlot := false
 		if slots := s.requestSlots; slots != nil {
 			select {
@@ -378,10 +379,11 @@ func (s *Server) serveConnection(baseCtx context.Context, state *connection) {
 		if requestSlot {
 			<-s.requestSlots
 		}
+		shouldClose := closeAfter || writer.closeAfter || closing
 		writer.reset("")
 		s.writerPool.Put(writer)
 		state.active.Store(false)
-		if closeAfter || closing {
+		if shouldClose {
 			return
 		}
 	}
@@ -719,11 +721,20 @@ func writeSimple(conn net.Conn, status int, body string, closeAfter bool, timeou
 }

 type responseWriter struct {
-	header        http.Header
-	status        int
-	wroteHeader   bool
-	body          bytes.Buffer
-	requestMethod string
+	header         http.Header
+	status         int
+	wroteHeader    bool
+	body           bytes.Buffer
+	requestMethod  string
+	protoMajor     int
+	protoMinor     int
+	conn           net.Conn
+	closeAfter     bool
+	writeTimeout   time.Duration
+	streaming      bool
+	chunked        bool
+	streamingClose bool
+	writeErr       error
 }

 func newResponseWriter(method string) *responseWriter {
@@ -742,6 +753,23 @@ func (w *responseWriter) reset(method string) {
 		w.body.Reset()
 	}
 	w.requestMethod = method
+	w.protoMajor = 0
+	w.protoMinor = 0
+	w.conn = nil
+	w.closeAfter = false
+	w.writeTimeout = 0
+	w.streaming = false
+	w.chunked = false
+	w.streamingClose = false
+	w.writeErr = nil
+}
+
+func (w *responseWriter) configure(request *http.Request, conn net.Conn, closeAfter bool, timeout time.Duration) {
+	w.protoMajor = request.ProtoMajor
+	w.protoMinor = request.ProtoMinor
+	w.conn = conn
+	w.closeAfter = closeAfter
+	w.writeTimeout = timeout
 }

 func (w *responseWriter) Header() http.Header { return w.header }
@@ -755,13 +783,121 @@ func (w *responseWriter) WriteHeader(status int) {
 }

 func (w *responseWriter) Write(data []byte) (int, error) {
+	if w.writeErr != nil {
+		return 0, w.writeErr
+	}
 	if !w.wroteHeader {
 		w.WriteHeader(http.StatusOK)
 	}
+	if w.streaming {
+		written, err := w.writeStream(data)
+		if err != nil {
+			w.writeErr = err
+		}
+		return written, err
+	}
 	return w.body.Write(data)
 }

+// Flush implements http.Flusher. Once a response is flushed, HTTP/1.1 uses
+// chunked framing unless the handler supplied Content-Length. HTTP/1.0 falls
+// back to close-delimited output because it has no chunked framing.
+func (w *responseWriter) Flush() {
+	if w.writeErr != nil || w.streaming || w.requestMethod == http.MethodHead || w.status == http.StatusNoContent ||
+		w.status == http.StatusResetContent || w.status == http.StatusNotModified || (w.status >= 100 && w.status < 200) {
+		return
+	}
+	if err := w.beginStream(); err != nil {
+		w.writeErr = err
+	}
+}
+
+func (w *responseWriter) beginStream() error {
+	if w.conn == nil {
+		return errors.New("go-omni/native: response stream is not connected")
+	}
+	if !w.wroteHeader {
+		w.status = http.StatusOK
+		w.wroteHeader = true
+	}
+	if w.requestMethod == http.MethodHead || w.status == http.StatusNoContent || w.status == http.StatusResetContent ||
+		w.status == http.StatusNotModified || (w.status >= 100 && w.status < 200) {
+		return nil
+	}
+	contentLength := w.header.Get("Content-Length")
+	w.chunked = contentLength == "" && w.protoMajor == 1 && w.protoMinor == 1
+	if w.chunked {
+		w.header.Set("Transfer-Encoding", "chunked")
+	} else if contentLength == "" {
+		w.streamingClose = true
+	}
+	closeAfter := w.closeAfter || w.streamingClose
+	if closeAfter {
+		w.header.Set("Connection", "close")
+	}
+	w.closeAfter = closeAfter
+	if err := w.writeHeaders(w.conn); err != nil {
+		return err
+	}
+	w.streaming = true
+	buffered := w.body.Bytes()
+	w.body.Reset()
+	if len(buffered) == 0 {
+		return nil
+	}
+	_, err := w.writeStream(buffered)
+	return err
+}
+
+func (w *responseWriter) writeStream(data []byte) (int, error) {
+	if len(data) == 0 {
+		return 0, nil
+	}
+	if w.conn == nil {
+		return 0, errors.New("go-omni/native: response stream is not connected")
+	}
+	if w.chunked {
+		var prefix [32]byte
+		framed := strconv.AppendInt(prefix[:0], int64(len(data)), 16)
+		framed = append(framed, '\r', '\n')
+		if err := w.writeBytes(framed); err != nil {
+			return 0, err
+		}
+		if err := w.writeBytes(data); err != nil {
+			return 0, err
+		}
+		if err := w.writeBytes([]byte("\r\n")); err != nil {
+			return 0, err
+		}
+		return len(data), nil
+	}
+	if err := w.writeBytes(data); err != nil {
+		return 0, err
+	}
+	return len(data), nil
+}
+
+func (w *responseWriter) writeBytes(data []byte) error {
+	if w.writeTimeout > 0 {
+		_ = w.conn.SetWriteDeadline(time.Now().Add(w.writeTimeout))
+	}
+	return writeAll(w.conn, data)
+}
+
 func (w *responseWriter) writeTo(conn net.Conn, closeAfter bool, timeout time.Duration) error {
+	if w.streaming {
+		if w.writeErr != nil {
+			_ = conn.SetWriteDeadline(time.Time{})
+			return w.writeErr
+		}
+		if w.chunked {
+			if err := w.writeBytes([]byte("0\r\n\r\n")); err != nil {
+				w.writeErr = err
+			}
+		}
+		_ = conn.SetWriteDeadline(time.Time{})
+		return w.writeErr
+	}
 	if timeout > 0 {
 		_ = conn.SetWriteDeadline(time.Now().Add(timeout))
 	}
@@ -775,7 +911,7 @@ func (w *responseWriter) writeTo(conn net.Conn, closeAfter bool, timeout time.Du
 	if noBody {
 		body = nil
 	}
-	if w.header.Get("Content-Length") == "" && w.status >= 200 && w.status != http.StatusNoContent && w.status != http.StatusNotModified && w.status != http.StatusResetContent {
+	if w.header.Get("Content-Length") == "" && w.header.Get("Transfer-Encoding") == "" && w.status >= 200 && w.status != http.StatusNoContent && w.status != http.StatusNotModified && w.status != http.StatusResetContent {
 		w.header.Set("Content-Length", strconv.Itoa(len(w.body.Bytes())))
 	}
 	if closeAfter {
@@ -808,7 +944,56 @@ func (w *responseWriter) writeTo(conn net.Conn, closeAfter bool, timeout time.Du
 	}
 	output.WriteString("\r\n")
 	output.Write(body)
-	_, err := conn.Write(output.Bytes())
+	err := writeAll(conn, output.Bytes())
 	_ = conn.SetWriteDeadline(time.Time{})
 	return err
 }
+
+func (w *responseWriter) writeHeaders(conn net.Conn) error {
+	if w.writeTimeout > 0 {
+		_ = conn.SetWriteDeadline(time.Now().Add(w.writeTimeout))
+	}
+	keys := make([]string, 0, len(w.header))
+	for key := range w.header {
+		keys = append(keys, key)
+	}
+	sort.Strings(keys)
+	var output bytes.Buffer
+	output.WriteString("HTTP/1.1 ")
+	output.WriteString(strconv.Itoa(w.status))
+	output.WriteByte(' ')
+	output.WriteString(http.StatusText(w.status))
+	output.WriteString("\r\n")
+	for _, key := range keys {
+		if !validToken(key) {
+			return errors.New("invalid response header name")
+		}
+		for _, value := range w.header.Values(key) {
+			if !validHeaderValue(value) {
+				return errors.New("invalid response header")
+			}
+			output.WriteString(key)
+			output.WriteString(": ")
+			output.WriteString(value)
+			output.WriteString("\r\n")
+		}
+	}
+	output.WriteString("\r\n")
+	return writeAll(conn, output.Bytes())
+}
+
+func writeAll(conn net.Conn, data []byte) error {
+	for len(data) > 0 {
+		written, err := conn.Write(data)
+		if written > 0 {
+			data = data[written:]
+		}
+		if err != nil {
+			return err
+		}
+		if written == 0 {
+			return io.ErrShortWrite
+		}
+	}
+	return nil
+}
--- a/transport/native/server_test.go
+++ b/transport/native/server_test.go
@@ -225,6 +225,60 @@ func TestServerPipelinedRequestsAndLargeBody(t *testing.T) {
 	}
 }

+func TestServerFlushesChunkedResponses(t *testing.T) {
+	flushed := make(chan struct{})
+	release := make(chan struct{})
+	application := omni.NewWithRouter(netrouter.New())
+	application.GET("/stream", func(ctx *omni.Context) error {
+		flusher, ok := ctx.Raw().(http.Flusher)
+		if !ok {
+			return fmt.Errorf("native response writer does not implement http.Flusher")
+		}
+		ctx.Raw().Header().Set("Content-Type", "text/plain")
+		_, _ = io.WriteString(ctx.Raw(), "one\n")
+		flusher.Flush()
+		close(flushed)
+		<-release
+		_, _ = io.WriteString(ctx.Raw(), "two\n")
+		return nil
+	})
+	server, listener, done := startServer(t, application)
+	defer stopServer(t, server, listener, done)
+
+	conn, err := net.Dial("tcp", listener.Addr().String())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer conn.Close()
+	reader := bufio.NewReader(conn)
+	_, _ = io.WriteString(conn, "GET /stream HTTP/1.1\r\nHost: test\r\nConnection: close\r\n\r\n")
+	response, err := http.ReadResponse(reader, &http.Request{Method: http.MethodGet})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if response.StatusCode != http.StatusOK || response.TransferEncoding == nil || response.TransferEncoding[0] != "chunked" {
+		response.Body.Close()
+		t.Fatalf("expected chunked response, status=%d transfer=%v", response.StatusCode, response.TransferEncoding)
+	}
+	select {
+	case <-flushed:
+	case <-time.After(time.Second):
+		response.Body.Close()
+		t.Fatal("handler did not flush response")
+	}
+	first := make([]byte, len("one\n"))
+	if _, err := io.ReadFull(response.Body, first); err != nil || string(first) != "one\n" {
+		response.Body.Close()
+		t.Fatalf("unexpected first stream chunk: %q error=%v", first, err)
+	}
+	close(release)
+	rest, err := io.ReadAll(response.Body)
+	response.Body.Close()
+	if err != nil || string(rest) != "two\n" {
+		t.Fatalf("unexpected remaining stream body: %q error=%v", rest, err)
+	}
+}
+
 func TestServerProtocolBoundaries(t *testing.T) {
 	application := omni.NewWithRouter(netrouter.New())
 	application.GET("/health", func(ctx *omni.Context) error { return ctx.OK("ok") })