Repository: cloudwego/hertz
Pull request #1343 (2025-04-29 09:54:17)

feat(sse): reader supports cancel stream
Also:
* ChunkedBodyWriter now respect io.Writer
* ChunkedBodyWriter returns err if it's finalized

#### What type of PR is this?
feat

<!--
Add one of the following kinds:

build: Changes that affect the build system or external dependencies (example scopes: gulp, broccoli, npm)
ci: Changes to our CI configuration files and scripts (example scopes: Travis, Circle, BrowserStack, SauceLabs)
docs: Documentation only changes
feat: A new feature
optimize: A new optimization
fix: A bug fix
perf: A code change that improves performance
refactor: A code change that neither fixes a bug nor adds a feature
style: Changes that do not affect the meaning of the code (white space, formatting, missing semi-colons, etc)
test: Adding missing tests or correcting existing tests
chore: Changes to the build process or auxiliary tools and libraries such as documentation generation
-->

#### Check the PR title.
<!--
The description of the title will be attached in Release Notes, 
so please describe it from user-oriented, what this PR does / why we need it.
Please check your PR title with the below requirements:
-->
- [ ] This PR title match the format: \<type\>(optional scope): \<description\>
- [ ] The description of this PR title is user-oriented and clear enough for others to understand.
- [ ] Attach the PR updating the user documentation if the current PR requires user awareness at the usage level. [User docs repo](https://github.com/cloudwego/cloudwego.github.io)


#### (Optional) Translate the PR title into Chinese.


#### (Optional) More detailed description for this PR(en: English/zh: Chinese).
<!--
Provide more detailed info for review(e.g., it's recommended to provide perf data if this is a perf type PR).
-->
en: reader supports cancel stream
zh(optional):

#### (Optional) Which issue(s) this PR fixes:
<!--
Automatically closes linked issue when PR is merged.
Eg: `Fixes #<issue number>`, or `Fixes (paste link of issue)`.
-->

#### (Optional) The PR that updates user documentation:
<!--
If the current PR requires user awareness at the usage level, please submit a PR to update user docs. [User docs repo](https://github.com/cloudwego/cloudwego.github.io)
-->

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/cloudwego/hertz/pull/1343?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego) Report
Attention: Patch coverage is `72.95082%` with `33 lines` in your changes missing coverage. Please review.
> Project coverage is 73.47%. Comparing base [(`57fbbf8`)](https://app.codecov.io/gh/cloudwego/hertz/commit/57fbbf8c1c61d3ebb7d729a771164e17af9913d2?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego) to head [(`3327cda`)](https://app.codecov.io/gh/cloudwego/hertz/commit/3327cda49175c7f78fa42b3da82054223e1f2d80?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego).

| [Files with missing lines](https://app.codecov.io/gh/cloudwego/hertz/pull/1343?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego) | Patch % | Lines |
|---|---|---|
| [pkg/protocol/http1/resp/writer.go](https://app.codecov.io/gh/cloudwego/hertz/pull/1343?src=pr&el=tree&filepath=pkg%2Fprotocol%2Fhttp1%2Fresp%2Fwriter.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL3Byb3RvY29sL2h0dHAxL3Jlc3Avd3JpdGVyLmdv) | 56.81% | [12 Missing and 7 partials :warning: ](https://app.codecov.io/gh/cloudwego/hertz/pull/1343?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego) |
| [pkg/protocol/http1/resp/response.go](https://app.codecov.io/gh/cloudwego/hertz/pull/1343?src=pr&el=tree&filepath=pkg%2Fprotocol%2Fhttp1%2Fresp%2Fresponse.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL3Byb3RvY29sL2h0dHAxL3Jlc3AvcmVzcG9uc2UuZ28=) | 55.55% | [8 Missing :warning: ](https://app.codecov.io/gh/cloudwego/hertz/pull/1343?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego) |
| [pkg/protocol/http1/ext/common.go](https://app.codecov.io/gh/cloudwego/hertz/pull/1343?src=pr&el=tree&filepath=pkg%2Fprotocol%2Fhttp1%2Fext%2Fcommon.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL3Byb3RvY29sL2h0dHAxL2V4dC9jb21tb24uZ28=) | 80.00% | [1 Missing and 2 partials :warning: ](https://app.codecov.io/gh/cloudwego/hertz/pull/1343?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego) |
| [pkg/protocol/sse/reader.go](https://app.codecov.io/gh/cloudwego/hertz/pull/1343?src=pr&el=tree&filepath=pkg%2Fprotocol%2Fsse%2Freader.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL3Byb3RvY29sL3NzZS9yZWFkZXIuZ28=) | 86.95% | [2 Missing and 1 partial :warning: ](https://app.codecov.io/gh/cloudwego/hertz/pull/1343?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego) |

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


```diff
@@             Coverage Diff             @@
##           develop    #1343      +/-   ##
===========================================
+ Coverage    73.45%   73.47%   +0.01%     
===========================================
  Files          144      144              
  Lines        16580    16639      +59     
===========================================
+ Hits         12179    12225      +46     
- Misses        3813     3824      +11     
- Partials       588      590       +2     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/cloudwego/hertz/pull/1343?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego).

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>

Patch:
diff --git a/internal/bytesconv/bytesconv.go b/internal/bytesconv/bytesconv.go
index e2d1ca5..898e222 100644
--- a/internal/bytesconv/bytesconv.go
+++ b/internal/bytesconv/bytesconv.go
@@ -43,7 +43,6 @@ package bytesconv
 
 import (
 	"net/http"
-	"sync"
 	"time"
 	"unsafe"
 
@@ -55,8 +54,6 @@ const (
 	lowerhex = "0123456789abcdef"
 )
 
-var hexIntBufPool sync.Pool
-
 func LowercaseBytes(b []byte) {
 	for i := 0; i < len(b); i++ {
 		p := &b[i]
@@ -89,29 +86,30 @@ func S2b(s string) (b []byte) {
 	return
 }
 
-func WriteHexInt(w network.Writer, n int) error {
-	if n < 0 {
-		panic("BUG: int must be positive")
+func EncodedIntHexLen(n uint64) int {
+	if n == 0 {
+		return 1
+	}
+	i := 0
+	for n > 0 {
+		i++
+		n >>= 4
 	}
+	return i
+}
 
-	v := hexIntBufPool.Get()
-	if v == nil {
-		v = make([]byte, maxHexIntChars+1)
+func AppendIntHex(b []byte, n uint64) []byte {
+	if n == 0 {
+		return append(b, '0')
 	}
-	buf := v.([]byte)
-	i := len(buf) - 1
-	for {
-		buf[i] = lowerhex[n&0xf]
-		n >>= 4
-		if n == 0 {
-			break
-		}
+	var tmp [16]byte // 64 / 4 = 16
+	i := len(tmp)
+	for n > 0 {
 		i--
+		tmp[i] = lowerhex[n&0xf]
+		n >>= 4
 	}
-	safeBuf, err := w.Malloc(maxHexIntChars + 1 - i)
-	copy(safeBuf, buf[i:])
-	hexIntBufPool.Put(v)
-	return err
+	return append(b, tmp[i:]...)
 }
 
 func ReadHexInt(r network.Reader) (int, error) {
diff --git a/pkg/network/writer.go b/pkg/network/writer.go
index ed6c6ac..483c96d 100644
--- a/pkg/network/writer.go
+++ b/pkg/network/writer.go
@@ -106,6 +106,8 @@ func NewWriter(w io.Writer) Writer {
 
 type ExtWriter interface {
 	io.Writer
+
+	// Flush sends data to peer immediately
 	Flush() error
 
 	// Finalize will be called by framework before the writer is released.
diff --git a/pkg/protocol/http1/ext/common.go b/pkg/protocol/http1/ext/common.go
index 2f2280a..90cd3b9 100644
--- a/pkg/protocol/http1/ext/common.go
+++ b/pkg/protocol/http1/ext/common.go
@@ -296,27 +296,31 @@ func round2(n int) int {
 	return 1 << x
 }
 
-func WriteChunk(w network.Writer, b []byte, withFlush bool) (err error) {
+func WriteChunk(w network.Writer, b []byte, withFlush bool) error {
 	n := len(b)
-	if err = bytesconv.WriteHexInt(w, n); err != nil {
-		return err
-	}
 
-	w.WriteBinary(bytestr.StrCRLF) //nolint:errcheck
-	if _, err = w.WriteBinary(b); err != nil {
+	sz := bytesconv.EncodedIntHexLen(uint64(n)) + 2 // len + CRLF
+	if n > 0 {
+		sz += n + 2 // data + CRLF
+	}
+	wb, err := w.Malloc(sz)
+	if err != nil {
 		return err
 	}
 
-	// If it is the end of chunk, write CRLF after writing trailer
+	wb = bytesconv.AppendIntHex(wb[:0], uint64(n))
+	wb = append(wb, '\r', '\n')
 	if n > 0 {
-		w.WriteBinary(bytestr.StrCRLF) //nolint:errcheck
+		wb = append(wb, b...)
+		wb = append(wb, '\r', '\n')
+	}
+	if len(wb) != sz {
+		panic("[BUG] len mismatch")
 	}
-
 	if !withFlush {
 		return nil
 	}
-	err = w.Flush()
-	return
+	return w.Flush()
 }
 
 func isOnlyCRLF(b []byte) bool {
diff --git a/pkg/protocol/http1/resp/header.go b/pkg/protocol/http1/resp/header.go
index ff12cb2..6d4e3b4 100644
--- a/pkg/protocol/http1/resp/header.go
+++ b/pkg/protocol/http1/resp/header.go
@@ -87,12 +87,14 @@ func ReadHeader(h *protocol.ResponseHeader, r network.Reader) error {
 
 // WriteHeader writes response header to w.
 func WriteHeader(h *protocol.ResponseHeader, w network.Writer) error {
+	// Data may become invalid after the next call of ResponseHeader.
+	// copy before WriteHeader returns
 	header := h.Header()
-	h.SetHeaderLength(len(header))
-	_, err := w.WriteBinary(header)
+	b, err := w.Malloc(len(header))
 	if err != nil {
 		return err
 	}
+	h.SetHeaderLength(copy(b, header))
 	return nil
 }
 
diff --git a/pkg/protocol/http1/resp/response.go b/pkg/protocol/http1/resp/response.go
index c55743f..6fa6e12 100644
--- a/pkg/protocol/http1/resp/response.go
+++ b/pkg/protocol/http1/resp/response.go
@@ -140,24 +140,49 @@ func ReadRespBody(resp *protocol.Response, r network.Reader, maxBodySize int) (e
 }
 
 type clientRespStream struct {
+	mu sync.Mutex
+
 	r             io.Reader
 	closeCallback func(shouldClose bool) error
 }
 
+// ForceClose closes underlying conn. It enables `Read` call to return instead of blocking.
+//
+// This method is ONLY used by hertz internally.
+// Normally, users call `Close` when the body is no longer used.
+func (c *clientRespStream) ForceClose() (err error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if c.closeCallback != nil {
+		err = c.closeCallback(true)
+		c.closeCallback = nil
+	}
+	// NOTE: DO NOT put back to pool here,
+	// user may still use clientRespStream and call Close() like `defer body.Close()`
+	return
+}
+
+// Close closes response stream gracefully.
+//
+// NOTE:
+// * Since Close() will put it back to pool, MUST ensure it only be called when no longer use.
+// * MUST NOT call Close() and `Read()` concurrently to avoid race issue
 func (c *clientRespStream) Close() (err error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
 	runtime.SetFinalizer(c, nil)
-	// If error happened in release, the connection may be in abnormal state.
+	// If error happened in ReleaseBodyStream, the connection may be in abnormal state.
 	// Close it in the callback in order to avoid other unexpected problems.
 	err = ext.ReleaseBodyStream(c.r)
-	shouldClose := false
-	if err != nil {
-		shouldClose = true
-		hlog.Warnf("connection will be closed instead of recycled because an error occurred during the stream body release: %s", err.Error())
-	}
 	if c.closeCallback != nil {
-		err = c.closeCallback(shouldClose)
+		if err != nil {
+			hlog.SystemLogger().Warnf("error occurred during the stream body close: %s", err)
+		}
+		err = c.closeCallback(err != nil)
 	}
-	c.reset()
+	c.r = nil
+	c.closeCallback = nil
+	clientRespStreamPool.Put(c)
 	return
 }
 
@@ -165,12 +190,6 @@ func (c *clientRespStream) Read(p []byte) (n int, err error) {
 	return c.r.Read(p)
 }
 
-func (c *clientRespStream) reset() {
-	c.closeCallback = nil
-	c.r = nil
-	clientRespStreamPool.Put(c)
-}
-
 var clientRespStreamPool = sync.Pool{
 	New: func() interface{} {
 		return &clientRespStream{}
diff --git a/pkg/protocol/http1/resp/writer.go b/pkg/protocol/http1/resp/writer.go
index 7b50cd4..02a3b7c 100644
--- a/pkg/protocol/http1/resp/writer.go
+++ b/pkg/protocol/http1/resp/writer.go
@@ -17,6 +17,7 @@
 package resp
 
 import (
+	"errors"
 	"runtime"
 	"sync"
 
@@ -36,32 +37,52 @@ func init() {
 }
 
 type chunkedBodyWriter struct {
-	sync.Once
-	finalizeErr error
+	r *protocol.Response
+	w network.Writer
+
+	err         error
+	finalized   bool
 	wroteHeader bool
-	r           *protocol.Response
-	w           network.Writer
 }
 
-// Write will encode chunked p before writing
-// It will only return the length of p and a nil error if the writing is successful or 0, error otherwise.
-//
-// NOTE: Write will use the user buffer to flush.
-// Before flush successfully, the buffer b should be valid.
+var errChunkedFinished = errors.New("chunked response is finished; no more data will be written.")
+
+// Write implements network.ExtWriter.Write / io.Writer.Write
 func (c *chunkedBodyWriter) Write(p []byte) (n int, err error) {
-	if !c.wroteHeader {
-		c.r.Header.SetContentLength(-1)
-		if err = WriteHeader(&c.r.Header, c.w); err != nil {
-			return
-		}
-		c.wroteHeader = true
+	if c.finalized {
+		return 0, errChunkedFinished
+	}
+	if c.err != nil {
+		return 0, c.err
+	}
+	if err := c.writeHeader(); err != nil {
+		return 0, err
 	}
-	if err = ext.WriteChunk(c.w, p, false); err != nil {
-		return
+	if err := c.writeChunk(p); err != nil {
+		return 0, err
 	}
 	return len(p), nil
 }
 
+func (c *chunkedBodyWriter) writeHeader() error {
+	if c.wroteHeader {
+		return nil
+	}
+	c.wroteHeader = true
+	c.r.Header.SetContentLength(-1)
+	if c.err = WriteHeader(&c.r.Header, c.w); c.err != nil {
+		return c.err
+	}
+	return nil
+}
+
+func (c *chunkedBodyWriter) writeChunk(b []byte) error {
+	if c.err = ext.WriteChunk(c.w, b, false); c.err != nil {
+		return c.err
+	}
+	return nil
+}
+
 func (c *chunkedBodyWriter) Flush() error {
 	return c.w.Flush()
 }
@@ -69,37 +90,36 @@ func (c *chunkedBodyWriter) Flush() error {
 // Finalize will write the ending chunk as well as trailer and flush the writer.
 // Warning: do not call this method by yourself, unless you know what you are doing.
 func (c *chunkedBodyWriter) Finalize() error {
-	c.Do(func() {
-		// in case no actual data from user
-		if !c.wroteHeader {
-			c.r.Header.SetContentLength(-1)
-			if c.finalizeErr = WriteHeader(&c.r.Header, c.w); c.finalizeErr != nil {
-				return
-			}
-			c.wroteHeader = true
-		}
-		c.finalizeErr = ext.WriteChunk(c.w, nil, true)
-		if c.finalizeErr != nil {
-			return
-		}
-		c.finalizeErr = ext.WriteTrailer(c.r.Header.Trailer(), c.w)
-	})
-	return c.finalizeErr
+	if c.finalized || c.err != nil {
+		return c.err
+	}
+	// zero-len chunk
+	if _, err := c.Write(nil); err != nil {
+		return err
+	}
+	c.finalized = true
+	// trailer which ends with \r\n
+	_, c.err = c.w.WriteBinary(c.r.Header.Trailer().Header())
+	if c.err == nil {
+		c.err = c.Flush()
+	}
+	return c.err
 }
 
 func (c *chunkedBodyWriter) release() {
 	c.r = nil
 	c.w = nil
-	c.finalizeErr = nil
+	c.err = nil
+	c.finalized = false
 	c.wroteHeader = false
 	chunkReaderPool.Put(c)
 }
 
+// NewChunkedBodyWriter creates a new chunked body writer.
 func NewChunkedBodyWriter(r *protocol.Response, w network.Writer) network.ExtWriter {
 	extWriter := chunkReaderPool.Get().(*chunkedBodyWriter)
 	extWriter.r = r
 	extWriter.w = w
-	extWriter.Once = sync.Once{}
 	runtime.SetFinalizer(extWriter, (*chunkedBodyWriter).release)
 	return extWriter
 }
diff --git a/pkg/protocol/response.go b/pkg/protocol/response.go
index 7aed602..3d8a80b 100644
--- a/pkg/protocol/response.go
+++ b/pkg/protocol/response.go
@@ -393,6 +393,10 @@ func (resp *Response) ConnectionClose() bool {
 	return resp.Header.ConnectionClose()
 }
 
+// CloseBodyStream tries call Close() of underlying body stream.
+//
+// NOTE:
+// * MUST NOT call CloseBodyStream() and BodyStream().Read() concurrently to avoid race issue.
 func (resp *Response) CloseBodyStream() error {
 	if resp.bodyStream == nil {
 		return nil
diff --git a/pkg/protocol/sse/reader.go b/pkg/protocol/sse/reader.go
index f44e734..b879fee 100644
--- a/pkg/protocol/sse/reader.go
+++ b/pkg/protocol/sse/reader.go
@@ -19,6 +19,7 @@ package sse
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"errors"
 	"io"
 	"strconv"
@@ -37,6 +38,7 @@ var errNotSSEContentType = errors.New("Content-Type returned by server is NOT te
 // It is used to parse the response body and extract individual events.
 type Reader struct {
 	resp   *protocol.Response
+	r      io.Reader
 	s      *bufio.Scanner
 	events int32
 
@@ -52,13 +54,18 @@ func NewReader(resp *protocol.Response) (*Reader, error) {
 	}
 	r := &Reader{resp: resp}
 	if resp.IsBodyStream() {
-		r.s = bufio.NewScanner(resp.BodyStream())
+		r.r = resp.BodyStream()
 	} else {
-		r.s = bufio.NewScanner(bytes.NewReader(resp.Body()))
+		r.r = bytes.NewReader(resp.Body())
 	}
+	r.s = bufio.NewScanner(r.r)
 	return r, nil
 }
 
+type forceCloseIf interface {
+	ForceClose() error // implemented by *clientRespStream
+}
+
 // ForEach iterates over all SSE events in the response body,
 // invoking the provided handler function for each event.
 //
@@ -66,15 +73,39 @@ func NewReader(resp *protocol.Response) (*Reader, error) {
 // Use (*Event).Clone to create a copy instead.
 //
 // Iteration stops if the handler returns an error or if reading fails.
+// `ctx` is used to cancel the iteration if ctx.Done() != nil
+//
 // It returns nil if all events are processed successfully.
-func (r *Reader) ForEach(f func(e *Event) error) error {
+func (r *Reader) ForEach(ctx context.Context, f func(e *Event) error) error {
+	if ctx.Done() != nil {
+		ch := make(chan struct{})
+		defer close(ch)
+		go func() {
+			select {
+			case <-ctx.Done():
+				// force close the underlying connection to release resource
+				// or r.Read may block until remote server ends
+				if s, ok := r.r.(forceCloseIf); ok {
+					s.ForceClose()
+				}
+			case <-ch:
+				return
+			}
+		}()
+	}
 	e := NewEvent()
 	defer e.Release()
 	for {
+		if err := ctx.Err(); err != nil {
+			return err
+		}
 		if err := r.Read(e); err != nil {
 			if err == io.EOF {
 				return nil
 			}
+			if er := ctx.Err(); er != nil {
+				err = er
+			}
 			return err
 		}
 		if err := f(e); err != nil {
@@ -175,6 +206,9 @@ func (r *Reader) Read(e *Event) error {
 }
 
 // Close closes the underlying response body.
+//
+// NOTE:
+// * MUST NOT call Close() and Read() / ForEach() concurrently to avoid race issue.
 func (r *Reader) Close() error {
 	return r.resp.CloseBodyStream()
 }
diff --git a/pkg/protocol/sse/writer.go b/pkg/protocol/sse/writer.go
index 9e6b8ba..511eecb 100644
--- a/pkg/protocol/sse/writer.go
+++ b/pkg/protocol/sse/writer.go
@@ -19,7 +19,6 @@ package sse
 import (
 	"bytes"
 	"errors"
-	"io"
 	"strconv"
 	"strings"
 	"sync"
@@ -27,19 +26,15 @@ import (
 	"github.com/cloudwego/hertz/internal/bytestr"
 	"github.com/cloudwego/hertz/pkg/app"
 	"github.com/cloudwego/hertz/pkg/common/bytebufferpool"
+	"github.com/cloudwego/hertz/pkg/network"
 	"github.com/cloudwego/hertz/pkg/protocol/http1/resp"
 )
 
-type writeFlusher interface {
-	io.Writer
-	Flush() error
-}
-
 // Writer represents a writer for Server-Sent Events (SSE).
 //
 // It is used to write individual events to the response body.
 type Writer struct {
-	w writeFlusher
+	w network.ExtWriter
 
 	mu sync.Mutex
 }
@@ -167,3 +162,9 @@ func (w *Writer) Write(e *Event) error {
 	}
 	return w.w.Flush()
 }
+
+func (w *Writer) Close() error {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	return w.w.Finalize()
+}


Test patch:
diff --git a/internal/bytesconv/bytesconv_32_test.go b/internal/bytesconv/bytesconv_32_test.go
index ea591f0..7b8ebdd 100644
--- a/internal/bytesconv/bytesconv_32_test.go
+++ b/internal/bytesconv/bytesconv_32_test.go
@@ -50,22 +50,6 @@ import (
 	"github.com/cloudwego/hertz/pkg/common/test/assert"
 )
 
-func TestWriteHexInt(t *testing.T) {
-	t.Parallel()
-
-	for _, v := range []struct {
-		s string
-		n int
-	}{
-		{"0", 0},
-		{"1", 1},
-		{"123", 0x123},
-		{"7fffffff", 0x7fffffff},
-	} {
-		testWriteHexInt(t, v.n, v.s)
-	}
-}
-
 func TestReadHexInt(t *testing.T) {
 	t.Parallel()
 
diff --git a/internal/bytesconv/bytesconv_64_test.go b/internal/bytesconv/bytesconv_64_test.go
index ec2dbaf..c724676 100644
--- a/internal/bytesconv/bytesconv_64_test.go
+++ b/internal/bytesconv/bytesconv_64_test.go
@@ -50,22 +50,6 @@ import (
 	"github.com/cloudwego/hertz/pkg/common/test/assert"
 )
 
-func TestWriteHexInt(t *testing.T) {
-	t.Parallel()
-
-	for _, v := range []struct {
-		s string
-		n int
-	}{
-		{"0", 0},
-		{"1", 1},
-		{"123", 0x123},
-		{"7fffffffffffffff", 0x7fffffffffffffff},
-	} {
-		testWriteHexInt(t, v.n, v.s)
-	}
-}
-
 func TestReadHexInt(t *testing.T) {
 	t.Parallel()
 
diff --git a/internal/bytesconv/bytesconv_test.go b/internal/bytesconv/bytesconv_test.go
index 316c917..b63f337 100644
--- a/internal/bytesconv/bytesconv_test.go
+++ b/internal/bytesconv/bytesconv_test.go
@@ -22,10 +22,8 @@ import (
 	"testing"
 	"time"
 
-	"github.com/cloudwego/hertz/pkg/common/bytebufferpool"
 	"github.com/cloudwego/hertz/pkg/common/test/assert"
 	"github.com/cloudwego/hertz/pkg/common/test/mock"
-	"github.com/cloudwego/hertz/pkg/network"
 )
 
 func TestAppendDate(t *testing.T) {
@@ -103,18 +101,36 @@ func TestS2b(t *testing.T) {
 	}
 }
 
-// common test function for 32bit and 64bit
-func testWriteHexInt(t *testing.T, n int, expectedS string) {
-	w := bytebufferpool.Get()
-	zw := network.NewWriter(w)
-	if err := WriteHexInt(zw, n); err != nil {
-		t.Errorf("unexpected error when writing hex %x: %v", n, err)
+func TestAppendIntHex(t *testing.T) {
+	testCases := []struct {
+		b        []byte
+		n        uint64
+		expected string
+	}{
+		{[]byte{}, 0, "0"},
+		{[]byte{}, 1, "1"},
+		{[]byte{}, 10, "a"},
+		{[]byte{}, 15, "f"},
+		{[]byte{}, 16, "10"},
+		{[]byte{}, 255, "ff"},
+		{[]byte{}, 256, "100"},
+		{[]byte{}, 123456789, "75bcd15"},
+		{[]byte{}, 0xffffffffffffffff, "ffffffffffffffff"},
+		{[]byte("pre-"), 255, "pre-ff"},
+		{[]byte("start"), 0, "start0"},
 	}
-	if err := zw.Flush(); err != nil {
-		t.Fatalf("unexpected error when flushing hex %x: %v", n, err)
+	for _, tc := range testCases {
+		result := AppendIntHex(tc.b, tc.n)
+		if string(result) != tc.expected {
+			t.Fatalf("AppendIntHex(%q, %d) = %q; want %q", tc.b, tc.n, result, tc.expected)
+		}
+
+		actualLen := EncodedIntHexLen(tc.n)
+		expectedLen := len(result) - len(tc.b)
+		if actualLen != expectedLen {
+			t.Fatalf("EncodedIntHexLen(%d) = %d; want %d", tc.n, actualLen, expectedLen)
+		}
 	}
-	s := B2s(w.B)
-	assert.DeepEqual(t, s, expectedS)
 }
 
 // common test function for 32bit and 64bit
diff --git a/pkg/common/test/mock/network.go b/pkg/common/test/mock/network.go
index b3ae54b..366e335 100644
--- a/pkg/common/test/mock/network.go
+++ b/pkg/common/test/mock/network.go
@@ -374,3 +374,36 @@ func (m *StreamConn) ReadBinary(n int) (p []byte, err error) {
 func DialerFun(addr string) (network.Conn, error) {
 	return NewConn(""), nil
 }
+
+type MockWriter struct {
+	w network.Writer
+
+	MockMalloc      func(n int) (buf []byte, err error)
+	MockWriteBinary func(b []byte) (n int, err error)
+	MockFlush       func() error
+}
+
+func NewMockWriter(w network.Writer) *MockWriter {
+	return &MockWriter{w: w}
+}
+
+func (m *MockWriter) Malloc(n int) (buf []byte, err error) {
+	if m.MockMalloc != nil {
+		return m.MockMalloc(n)
+	}
+	return m.w.Malloc(n)
+}
+
+func (m *MockWriter) WriteBinary(b []byte) (n int, err error) {
+	if m.MockWriteBinary != nil {
+		return m.MockWriteBinary(b)
+	}
+	return m.w.WriteBinary(b)
+}
+
+func (m *MockWriter) Flush() error {
+	if m.MockFlush != nil {
+		return m.MockFlush()
+	}
+	return m.w.Flush()
+}
diff --git a/pkg/protocol/http1/resp/writer_test.go b/pkg/protocol/http1/resp/writer_test.go
index b57281a..80bec2f 100644
--- a/pkg/protocol/http1/resp/writer_test.go
+++ b/pkg/protocol/http1/resp/writer_test.go
@@ -17,10 +17,10 @@
 package resp
 
 import (
+	"errors"
 	"strings"
 	"testing"
 
-	"github.com/cloudwego/hertz/internal/bytestr"
 	"github.com/cloudwego/hertz/pkg/common/test/assert"
 	"github.com/cloudwego/hertz/pkg/common/test/mock"
 	"github.com/cloudwego/hertz/pkg/protocol"
@@ -28,39 +28,52 @@ import (
 
 func TestNewChunkedBodyWriter(t *testing.T) {
 	response := protocol.AcquireResponse()
-	mockConn := mock.NewConn("")
-	w := NewChunkedBodyWriter(response, mockConn)
-	w.Write([]byte("hello"))
-	w.Flush()
-	out, _ := mockConn.WriterRecorder().ReadBinary(mockConn.WriterRecorder().WroteLen())
-	assert.True(t, strings.Contains(string(out), "Transfer-Encoding: chunked"))
-	assert.True(t, strings.Contains(string(out), "5"+string(bytestr.StrCRLF)+"hello"))
-	assert.False(t, strings.Contains(string(out), "0"+string(bytestr.StrCRLF)+string(bytestr.StrCRLF)))
-}
+	defer protocol.ReleaseResponse(response)
 
-func TestNewChunkedBodyWriter1(t *testing.T) {
-	response := protocol.AcquireResponse()
 	mockConn := mock.NewConn("")
 	w := NewChunkedBodyWriter(response, mockConn)
-	w.Write([]byte("hello"))
-	w.Flush()
-	w.Finalize()
-	w.Flush()
-	out, _ := mockConn.WriterRecorder().ReadBinary(mockConn.WriterRecorder().WroteLen())
-	assert.True(t, strings.Contains(string(out), "Transfer-Encoding: chunked"))
-	assert.True(t, strings.Contains(string(out), "5"+string(bytestr.StrCRLF)+"hello"))
-	assert.True(t, strings.Contains(string(out), "0"+string(bytestr.StrCRLF)+string(bytestr.StrCRLF)))
+	_, _ = w.Write([]byte("hello"))
+	assert.Nil(t, w.Flush())
+	out, _ := mockConn.WriterRecorder().Peek(mockConn.WriterRecorder().WroteLen())
+	resp := string(out)
+	assert.True(t, strings.Contains(resp, "Transfer-Encoding: chunked"))
+	assert.True(t, strings.HasSuffix(resp, "5\r\nhello\r\n"))
+
+	// Finalize adds 0\r\n\r\n
+	assert.Nil(t, w.Finalize())
+	assert.Nil(t, w.Finalize()) // noop
+	out, _ = mockConn.WriterRecorder().Peek(mockConn.WriterRecorder().WroteLen())
+	resp = string(out)
+	assert.True(t, strings.HasSuffix(resp, "5\r\nhello\r\n0\r\n\r\n"))
+
+	_, err := w.Write([]byte("world"))
+	assert.True(t, err == errChunkedFinished)
 }
 
-func TestNewChunkedBodyWriterNoData(t *testing.T) {
+func TestNewChunkedBodyWriter_Err(t *testing.T) {
 	response := protocol.AcquireResponse()
-	response.Header.Set("Foo", "Bar")
-	mockConn := mock.NewConn("")
-	w := NewChunkedBodyWriter(response, mockConn)
-	w.Finalize()
-	w.Flush()
-	out, _ := mockConn.WriterRecorder().ReadBinary(mockConn.WriterRecorder().WroteLen())
-	assert.True(t, strings.Contains(string(out), "Transfer-Encoding: chunked"))
-	assert.True(t, strings.Contains(string(out), "Foo: Bar"))
-	assert.True(t, strings.Contains(string(out), "0"+string(bytestr.StrCRLF)+string(bytestr.StrCRLF)))
+	defer protocol.ReleaseResponse(response)
+
+	mw := mock.NewMockWriter(nil)
+	w := NewChunkedBodyWriter(response, mw)
+
+	expectErr := errors.New("mock malloc err")
+
+	mw.MockMalloc = func(n int) ([]byte, error) {
+		return nil, expectErr
+	}
+	_, err := w.Write([]byte("hello"))
+	assert.True(t, err == expectErr)
+
+	mw.MockMalloc = nil
+	_, err = w.Write([]byte("world"))
+	assert.True(t, err == expectErr) // next call will return last err
+
+	w = NewChunkedBodyWriter(response, mw)
+
+	mw.MockMalloc = func(n int) ([]byte, error) {
+		return nil, expectErr
+	}
+	err = w.Finalize()
+	assert.True(t, err == expectErr)
 }
diff --git a/pkg/protocol/sse/example_test.go b/pkg/protocol/sse/example_test.go
index 7f83e76..44a58aa 100644
--- a/pkg/protocol/sse/example_test.go
+++ b/pkg/protocol/sse/example_test.go
@@ -42,11 +42,13 @@ func Example() {
 			w.WriteEvent(fmt.Sprintf("id-%d", i), "message", []byte("hello\n\nworld"))
 			time.Sleep(10 * time.Millisecond)
 		}
+		// [optional] it writes 0\r\n\r\n to indicate the end of chunked response
+		// hertz will do it after handler returns
+		w.Close()
 	})
 	go engine.Run()
 	defer engine.Close()
-	// Wait for server to start
-	time.Sleep(20 * time.Millisecond)
+	time.Sleep(20 * time.Millisecond) // wait for server to start
 	opt.Addr = testutils.GetListenerAddr(engine)
 
 	// --- SSE Client ---
@@ -55,7 +57,11 @@ func Example() {
 	req.SetRequestURI("http://" + opt.Addr + "/")
 	req.SetMethod("GET")
 	req.SetHeader(LastEventIDHeader, "id-0")
-	AddAcceptMIME(req) // optional for most SSE servers
+
+	// adds `text/event-stream` to http `Accept` header
+	// may required for some Model Context Protocol(MCP) servers
+	AddAcceptMIME(req)
+
 	if err := c.Do(context.Background(), req, resp); err != nil {
 		panic(err)
 	}
@@ -63,7 +69,15 @@ func Example() {
 	if err != nil {
 		panic(err)
 	}
-	err = r.ForEach(func(e *Event) error {
+	defer r.Close()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	go func() {
+		time.Sleep(200 * time.Millisecond)
+		// cancel can be used to force ForEach returns by closing the remote connection
+		_ = cancel
+	}()
+	err = r.ForEach(ctx, func(e *Event) error {
 		println("Event:", e.String())
 		return nil
 	})
diff --git a/pkg/protocol/sse/reader_test.go b/pkg/protocol/sse/reader_test.go
index 9b62496..9de6ebf 100644
--- a/pkg/protocol/sse/reader_test.go
+++ b/pkg/protocol/sse/reader_test.go
@@ -18,6 +18,8 @@ package sse
 
 import (
 	"bytes"
+	"context"
+	"errors"
 	"io"
 	"strings"
 	"testing"
@@ -277,6 +279,19 @@ func TestReader_ReadEvent_WithBodyStream(t *testing.T) {
 	e.Release()
 }
 
+type mockReadForceClose struct {
+	readFunc  func(b []byte) (int, error)
+	closeFunc func() error
+}
+
+func (m *mockReadForceClose) Read(b []byte) (int, error) {
+	return m.readFunc(b)
+}
+
+func (m *mockReadForceClose) ForceClose() error {
+	return m.closeFunc()
+}
+
 func TestReader_ReadEvent_Error(t *testing.T) {
 	// Create a reader that will return an error
 	errReader := &bytes.Reader{}
@@ -296,3 +311,35 @@ func TestReader_ReadEvent_Error(t *testing.T) {
 
 	e.Release()
 }
+
+func TestReader_ForEach(t *testing.T) {
+	// mock Read & ForceClose
+	mr := &mockReadForceClose{}
+	ch := make(chan error, 1)
+	defer close(ch)
+	mr.readFunc = func(b []byte) (int, error) {
+		return 0, <-ch
+	}
+	mr.closeFunc = func() error {
+		ch <- errors.New("closed")
+		return nil
+	}
+
+	// create protocol.Response
+	resp := &protocol.Response{}
+	resp.Header.SetContentType(string(bytestr.MIMETextEventStream))
+	resp.SetBodyStream(mr, -1)
+	r, err := NewReader(resp)
+	assert.Assert(t, err == nil)
+
+	// test ForEach with context
+	ctx, cancel := context.WithCancel(context.Background())
+	go func() { // cancel after 50ms
+		time.Sleep(50 * time.Millisecond)
+		cancel()
+	}()
+	err = r.ForEach(ctx, func(e *Event) error {
+		panic("must not called")
+	})
+	assert.Assert(t, err == ctx.Err())
+}
diff --git a/pkg/protocol/sse/writer_test.go b/pkg/protocol/sse/writer_test.go
index 24e2aa8..3c2c693 100644
--- a/pkg/protocol/sse/writer_test.go
+++ b/pkg/protocol/sse/writer_test.go
@@ -32,6 +32,7 @@ type mw struct {
 	flushCalled bool
 	writeErr    error
 	flushErr    error
+	finalizeErr error
 }
 
 func (m *mw) Write(p []byte) (n int, err error) {
@@ -50,6 +51,10 @@ func (m *mw) String() string {
 	return m.buf.String()
 }
 
+func (m *mw) Finalize() error {
+	return m.finalizeErr
+}
+
 func TestWriter_WriteEvent(t *testing.T) {
 	tests := []struct {
 		name      string
@@ -220,3 +225,27 @@ func TestWriter_WriteComment(t *testing.T) {
 	assert.Assert(t, err == nil)
 	assert.DeepEqual(t, ":keep-alive\n", m.String())
 }
+
+func TestWriter_Close(t *testing.T) {
+	// Create a mock writeFlusher
+	m := &mw{}
+
+	// Create a Writer with the mock
+	w := &Writer{w: m}
+
+	// Test Close method
+	err := w.Close()
+
+	// Verify no error occurred
+	assert.Nil(t, err)
+
+	// Set an error to be returned by Finalize
+	expectedErr := errors.New("finalize error")
+	m.finalizeErr = expectedErr
+
+	// Test Close method with error
+	err = w.Close()
+
+	// Verify the error is propagated
+	assert.DeepEqual(t, expectedErr, err)
+}