Repository: cloudwego/hertz
Pull request #610 (2023-02-13 12:09:56)

feat(http1): support flush chunk in handler
#### What type of PR is this?
feat

#### Check the PR title.
- [x] This PR title match the format: \<type\>(optional scope): \<description\>
- [x] The description of this PR title is user-oriented and clear enough for others to understand.


#### (Optional) Translate the PR title into Chinese.
支持在handler中按chunk flush数据

#### (Optional) More detail description for this PR(en: English/zh: Chinese).
<!--
Provide more detailed info for review. If it is a perf type PR, perf data is suggested to give.
-->
en: 
1. Support hijack for writer of response
2. Provide a impl for chunk-flushing writer for http/1.1

zh(optional):
1. response 支持writer劫持
2. 提供一个支持按chunk块flush的writer实现

#### Which issue(s) this PR fixes:
None


codecov[bot]: # [Codecov](https://codecov.io/gh/cloudwego/hertz/pull/610?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego) Report
Base: **72.12**% // Head: **71.88**% // Decreases project coverage by **`-0.24%`** :warning:
> Coverage data is based on head [(`7a255cf`)](https://codecov.io/gh/cloudwego/hertz/pull/610?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego) compared to base [(`e4cc656`)](https://codecov.io/gh/cloudwego/hertz/commit/e4cc65689777ea5fe6101a6e6037e1b57944c4c3?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego).
> Patch coverage: 8.10% of modified lines in pull request are covered.

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


```diff
@@             Coverage Diff             @@
##           develop     #610      +/-   ##
===========================================
- Coverage    72.12%   71.88%   -0.24%     
===========================================
  Files           94       95       +1     
  Lines         9234     9266      +32     
===========================================
+ Hits          6660     6661       +1     
- Misses        2155     2183      +28     
- Partials       419      422       +3     
```


| [Impacted Files](https://codecov.io/gh/cloudwego/hertz/pull/610?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego) | Coverage Δ | |
|---|---|---|
| [pkg/app/fs.go](https://codecov.io/gh/cloudwego/hertz/pull/610?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL2FwcC9mcy5nbw==) | `68.65% <ø> (ø)` | |
| [pkg/network/writer.go](https://codecov.io/gh/cloudwego/hertz/pull/610?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL25ldHdvcmsvd3JpdGVyLmdv) | `95.65% <ø> (ø)` | |
| [pkg/protocol/http1/ext/common.go](https://codecov.io/gh/cloudwego/hertz/pull/610?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL3Byb3RvY29sL2h0dHAxL2V4dC9jb21tb24uZ28=) | `21.13% <0.00%> (ø)` | |
| [pkg/protocol/http1/resp/writer.go](https://codecov.io/gh/cloudwego/hertz/pull/610?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL3Byb3RvY29sL2h0dHAxL3Jlc3Avd3JpdGVyLmdv) | `0.00% <0.00%> (ø)` | |
| [pkg/protocol/http1/server.go](https://codecov.io/gh/cloudwego/hertz/pull/610?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL3Byb3RvY29sL2h0dHAxL3NlcnZlci5nbw==) | `61.72% <0.00%> (-0.60%)` | :arrow_down: |
| [pkg/protocol/response.go](https://codecov.io/gh/cloudwego/hertz/pull/610?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL3Byb3RvY29sL3Jlc3BvbnNlLmdv) | `79.01% <23.07%> (-5.10%)` | :arrow_down: |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego)

</details>

[:umbrella: View full report at Codecov](https://codecov.io/gh/cloudwego/hertz/pull/610?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego).


Patch:
diff --git a/cmd/hz/app/app.go b/cmd/hz/app/app.go
index abf7c07..35a2e29 100644
--- a/cmd/hz/app/app.go
+++ b/cmd/hz/app/app.go
@@ -161,6 +161,7 @@ func Init() *cli.App {
 	handlerDirFlag := cli.StringFlag{Name: "handler_dir", Usage: "Specify the handler relative path (based on \"out_dir\").", Destination: &globalArgs.HandlerDir}
 	modelDirFlag := cli.StringFlag{Name: "model_dir", Usage: "Specify the model relative path (based on \"out_dir\").", Destination: &globalArgs.ModelDir}
 	routerDirFlag := cli.StringFlag{Name: "router_dir", Usage: "Specify the router relative path (based on \"out_dir\").", Destination: &globalArgs.RouterDir}
+	useFlag := cli.StringFlag{Name: "use", Usage: "Specify the model package to import for handler.", Destination: &globalArgs.Use}
 	baseDomainFlag := cli.StringFlag{Name: "base_domain", Usage: "Specify the request domain.", Destination: &globalArgs.BaseDomain}
 	clientDirFlag := cli.StringFlag{Name: "client_dir", Usage: "Specify the client path. If not specified, IDL generated path is used for 'client' command; no client code is generated for 'new' command", Destination: &globalArgs.ClientDir}
 
@@ -209,6 +210,7 @@ func Init() *cli.App {
 				&modelDirFlag,
 				&routerDirFlag,
 				&clientDirFlag,
+				&useFlag,
 
 				&includesFlag,
 				&thriftOptionsFlag,
@@ -267,6 +269,7 @@ func Init() *cli.App {
 				&moduleFlag,
 				&outDirFlag,
 				&modelDirFlag,
+				&useFlag,
 
 				&includesFlag,
 				&thriftOptionsFlag,
@@ -290,6 +293,7 @@ func Init() *cli.App {
 				&baseDomainFlag,
 				&modelDirFlag,
 				&clientDirFlag,
+				&useFlag,
 
 				&includesFlag,
 				&thriftOptionsFlag,
@@ -397,7 +401,10 @@ func TriggerPlugin(args *config.Argument) error {
 	logs.Debugf("begin to trigger plugin, compiler: %s, idl_paths: %v", compiler, args.IdlPaths)
 	buf, err := cmd.CombinedOutput()
 	if err != nil {
-		return fmt.Errorf("plugin %s_gen_hertz returns error: %v, cause:\n%v", compiler, err, string(buf))
+		out := strings.TrimSpace(string(buf))
+		if !strings.HasSuffix(out, meta.TheUseOptionMessage) {
+			return fmt.Errorf("plugin %s_gen_hertz returns error: %v, cause:\n%v", compiler, err, string(buf))
+		}
 	}
 
 	// If len(buf) != 0, the plugin returned the log.
diff --git a/cmd/hz/config/argument.go b/cmd/hz/config/argument.go
index 4ed71f5..994e785 100644
--- a/cmd/hz/config/argument.go
+++ b/cmd/hz/config/argument.go
@@ -52,6 +52,7 @@ type Argument struct {
 	Gomod       string
 	Gopkg       string // $GOPATH/src/{{gopkg}}
 	ServiceName string // service name
+	Use         string
 
 	JSONEnumStr          bool
 	UnsetOmitempty       bool
diff --git a/cmd/hz/generator/client.go b/cmd/hz/generator/client.go
index bbc3b5c..77470a3 100644
--- a/cmd/hz/generator/client.go
+++ b/cmd/hz/generator/client.go
@@ -18,6 +18,7 @@ package generator
 
 import (
 	"path/filepath"
+	"strings"
 
 	"github.com/cloudwego/hertz/cmd/hz/generator/model"
 	"github.com/cloudwego/hertz/cmd/hz/util"
@@ -78,6 +79,15 @@ func (pkgGen *HttpPackageGenerator) genClient(pkg *HttpPackage, clientDir string
 				client.Imports[mm.PackageName] = mm
 			}
 		}
+		if len(pkgGen.UseDir) != 0 {
+			oldModelDir := filepath.Clean(filepath.Join(pkgGen.ProjPackage, pkgGen.ModelDir))
+			newModelDir := filepath.Clean(pkgGen.UseDir)
+			for _, m := range client.ClientMethods {
+				for _, mm := range m.Models {
+					mm.Package = strings.Replace(mm.Package, oldModelDir, newModelDir, 1)
+				}
+			}
+		}
 		err = pkgGen.TemplateGenerator.Generate(client, idlClientName, client.FilePath, false)
 		if err != nil {
 			return err
diff --git a/cmd/hz/generator/handler.go b/cmd/hz/generator/handler.go
index aa76979..a39e0fe 100644
--- a/cmd/hz/generator/handler.go
+++ b/cmd/hz/generator/handler.go
@@ -130,6 +130,16 @@ func (pkgGen *HttpPackageGenerator) processHandler(handler *Handler, root *Route
 		}
 	}
 
+	if len(pkgGen.UseDir) != 0 {
+		oldModelDir := filepath.Clean(filepath.Join(pkgGen.ProjPackage, pkgGen.ModelDir))
+		newModelDir := filepath.Clean(pkgGen.UseDir)
+		for _, m := range handler.Methods {
+			for _, mm := range m.Models {
+				mm.Package = strings.Replace(mm.Package, oldModelDir, newModelDir, 1)
+			}
+		}
+	}
+
 	handler.Format()
 	return nil
 }
diff --git a/cmd/hz/generator/package.go b/cmd/hz/generator/package.go
index a0f44d4..b1d89e4 100644
--- a/cmd/hz/generator/package.go
+++ b/cmd/hz/generator/package.go
@@ -55,6 +55,7 @@ type HttpPackageGenerator struct {
 	HandlerDir      string
 	RouterDir       string
 	ModelDir        string
+	UseDir          string
 	ClientDir       string
 	IdlClientDir    string
 	NeedModel       bool
diff --git a/cmd/hz/meta/const.go b/cmd/hz/meta/const.go
index 2e4d82f..5c27962 100644
--- a/cmd/hz/meta/const.go
+++ b/cmd/hz/meta/const.go
@@ -81,3 +81,6 @@ const (
 const (
 	SetBodyParam = "setBodyParam(req).\n"
 )
+
+// TheUseOptionMessage indicates that the generating of 'model code' is aborted due to the -use option for thrift IDL.
+const TheUseOptionMessage = "'model code' is not generated due to the '-use' option"
diff --git a/cmd/hz/protobuf/plugin.go b/cmd/hz/protobuf/plugin.go
index 0e65a95..6f43382 100644
--- a/cmd/hz/protobuf/plugin.go
+++ b/cmd/hz/protobuf/plugin.go
@@ -74,6 +74,7 @@ type Plugin struct {
 	Recursive    bool
 	OutDir       string
 	ModelDir     string
+	UseDir       string
 	IdlClientDir string
 	PkgMap       map[string]string
 	logger       *logs.StdLogger
@@ -165,6 +166,7 @@ func (plugin *Plugin) parseArgs(param string) (*config.Argument, error) {
 	}
 	plugin.OutDir = args.OutDir
 	plugin.PkgMap = args.OptPkgMap
+	plugin.UseDir = args.Use
 	return args, nil
 }
 
@@ -349,7 +351,10 @@ func (plugin *Plugin) GenerateFile(gen *protogen.Plugin, f *protogen.File) error
 	}
 	f.GeneratedFilenamePrefix = filepath.Join(util.ImportToPath(impt, ""), util.BaseName(f.Proto.GetName(), ".proto"))
 	f.Generate = true
-
+	// if use third-party model, no model code is generated within the project
+	if len(plugin.UseDir) != 0 {
+		return nil
+	}
 	file, err := generateFile(gen, f)
 	if err != nil || file == nil {
 		return fmt.Errorf("generate file %s failed: %s", f.Proto.GetName(), err.Error())
@@ -586,9 +591,11 @@ func (plugin *Plugin) genHttpPackage(ast *descriptorpb.FileDescriptorProto, deps
 		HandlerDir: handlerDir,
 		RouterDir:  routerDir,
 		ModelDir:   modelDir,
+		UseDir:     args.Use,
 		ClientDir:  clientDir,
 		TemplateGenerator: generator.TemplateGenerator{
 			OutputDir: args.OutDir,
+			Excludes:  args.Excludes,
 		},
 		ProjPackage:     pkg,
 		Options:         options,
diff --git a/cmd/hz/thrift/plugin.go b/cmd/hz/thrift/plugin.go
index 44aa475..2d30233 100644
--- a/cmd/hz/thrift/plugin.go
+++ b/cmd/hz/thrift/plugin.go
@@ -17,6 +17,7 @@
 package thrift
 
 import (
+	"errors"
 	"fmt"
 	"io/ioutil"
 	"os"
@@ -128,9 +129,11 @@ func (plugin *Plugin) Run() int {
 		HandlerDir: handlerDir,
 		RouterDir:  routerDir,
 		ModelDir:   modelDir,
+		UseDir:     args.Use,
 		ClientDir:  clientDir,
 		TemplateGenerator: generator.TemplateGenerator{
 			OutputDir: args.OutDir,
+			Excludes:  args.Excludes,
 		},
 		ProjPackage:     pkg,
 		Options:         options,
@@ -149,6 +152,20 @@ func (plugin *Plugin) Run() int {
 		logs.Errorf("generate package failed: %s", err.Error())
 		return meta.PluginError
 	}
+	if len(args.Use) != 0 {
+		err = sg.Persist()
+		if err != nil {
+			logs.Errorf("persist file failed within '-use' option: %s", err.Error())
+			return meta.PluginError
+		}
+		res := thriftgo_plugin.BuildErrorResponse(errors.New(meta.TheUseOptionMessage).Error())
+		err = plugin.response(res)
+		if err != nil {
+			logs.Errorf("response failed: %s", err.Error())
+			return meta.PluginError
+		}
+		return 0
+	}
 	files, err := sg.GetFormatAndExcludedFiles()
 	if err != nil {
 		logs.Errorf("format file failed: %s", err.Error())
@@ -159,7 +176,7 @@ func (plugin *Plugin) Run() int {
 		logs.Errorf("get response failed: %s", err.Error())
 		return meta.PluginError
 	}
-	plugin.response(res)
+	err = plugin.response(res)
 	if err != nil {
 		logs.Errorf("response failed: %s", err.Error())
 		return meta.PluginError
diff --git a/pkg/app/context.go b/pkg/app/context.go
index e86cc4b..3ed3f19 100644
--- a/pkg/app/context.go
+++ b/pkg/app/context.go
@@ -211,6 +211,15 @@ type RequestContext struct {
 	formValueFunc FormValueFunc
 }
 
+// Flush is the shortcut for ctx.Response.GetHijackWriter().Flush().
+// Will return nil if the response writer is not hijacked.
+func (ctx *RequestContext) Flush() error {
+	if ctx.Response.GetHijackWriter() == nil {
+		return nil
+	}
+	return ctx.Response.GetHijackWriter().Flush()
+}
+
 func (ctx *RequestContext) SetClientIPFunc(f ClientIP) {
 	ctx.clientIPFunc = f
 }
diff --git a/pkg/app/fs.go b/pkg/app/fs.go
index 72aad87..5bad7cc 100644
--- a/pkg/app/fs.go
+++ b/pkg/app/fs.go
@@ -276,8 +276,6 @@ func (r *fsSmallFileReader) WriteTo(w io.Writer) (int64, error) {
 // Directory contents is returned if path points to directory.
 //
 // Use ServeFileUncompressed is you don't need serving compressed file contents.
-//
-// See also RequestCtx.SendFile.
 func ServeFile(ctx *RequestContext, path string) {
 	rootFSOnce.Do(func() {
 		rootFSHandler = rootFS.NewRequestHandler()
@@ -1204,8 +1202,6 @@ func stripLeadingSlashes(path []byte, stripSlashes int) []byte {
 //
 // ServeFile may be used for saving network traffic when serving files
 // with good compression ratio.
-//
-// See also RequestCtx.SendFile.
 func ServeFileUncompressed(ctx *RequestContext, path string) {
 	ctx.Request.Header.DelBytes(bytestr.StrAcceptEncoding)
 	ServeFile(ctx, path)
diff --git a/pkg/network/writer.go b/pkg/network/writer.go
index 0b3a405..ed6c6ac 100644
--- a/pkg/network/writer.go
+++ b/pkg/network/writer.go
@@ -103,3 +103,12 @@ func NewWriter(w io.Writer) Writer {
 		w: w,
 	}
 }
+
+type ExtWriter interface {
+	io.Writer
+	Flush() error
+
+	// Finalize will be called by framework before the writer is released.
+	// Implementations must guarantee that Finalize is safe for multiple calls.
+	Finalize() error
+}
diff --git a/pkg/protocol/header.go b/pkg/protocol/header.go
index 693ae47..d961418 100644
--- a/pkg/protocol/header.go
+++ b/pkg/protocol/header.go
@@ -87,7 +87,7 @@ type RequestHeader struct {
 
 	h       []argsKV
 	bufKV   argsKV
-	trailer Trailer
+	trailer *Trailer
 
 	cookies []argsKV
 
@@ -127,7 +127,7 @@ type ResponseHeader struct {
 
 	h       []argsKV
 	bufKV   argsKV
-	trailer Trailer
+	trailer *Trailer
 
 	cookies []argsKV
 
@@ -1788,10 +1788,16 @@ func (h *RequestHeader) setSpecialHeader(key, value []byte) bool {
 
 // Trailer returns the Trailer of HTTP Header.
 func (h *ResponseHeader) Trailer() *Trailer {
-	return &h.trailer
+	if h.trailer == nil {
+		h.trailer = new(Trailer)
+	}
+	return h.trailer
 }
 
 // Trailer returns the Trailer of HTTP Header.
 func (h *RequestHeader) Trailer() *Trailer {
-	return &h.trailer
+	if h.trailer == nil {
+		h.trailer = new(Trailer)
+	}
+	return h.trailer
 }
diff --git a/pkg/protocol/http1/ext/common.go b/pkg/protocol/http1/ext/common.go
index 94526b7..b31b40c 100644
--- a/pkg/protocol/http1/ext/common.go
+++ b/pkg/protocol/http1/ext/common.go
@@ -118,14 +118,14 @@ func WriteBodyChunked(w network.Writer, r io.Reader) error {
 				panic("BUG: io.Reader returned 0, nil")
 			}
 			if err == io.EOF {
-				if err = writeChunk(w, buf[:0]); err != nil {
+				if err = WriteChunk(w, buf[:0], true); err != nil {
 					break
 				}
 				err = nil
 			}
 			break
 		}
-		if err = writeChunk(w, buf[:n]); err != nil {
+		if err = WriteChunk(w, buf[:n], true); err != nil {
 			break
 		}
 	}
@@ -287,7 +287,7 @@ func round2(n int) int {
 	return 1 << x
 }
 
-func writeChunk(w network.Writer, b []byte) (err error) {
+func WriteChunk(w network.Writer, b []byte, withFlush bool) (err error) {
 	n := len(b)
 	if err = bytesconv.WriteHexInt(w, n); err != nil {
 		return err
@@ -303,6 +303,9 @@ func writeChunk(w network.Writer, b []byte) (err error) {
 		w.WriteBinary(bytestr.StrCRLF) //nolint:errcheck
 	}
 
+	if !withFlush {
+		return nil
+	}
 	err = w.Flush()
 	return
 }
diff --git a/pkg/protocol/http1/resp/writer.go b/pkg/protocol/http1/resp/writer.go
new file mode 100644
index 0000000..745ab4c
--- /dev/null
+++ b/pkg/protocol/http1/resp/writer.go
@@ -0,0 +1,76 @@
+/*
+ * Copyright 2023 CloudWeGo Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package resp
+
+import (
+	"sync"
+
+	"github.com/cloudwego/hertz/pkg/network"
+	"github.com/cloudwego/hertz/pkg/protocol"
+	"github.com/cloudwego/hertz/pkg/protocol/http1/ext"
+)
+
+type chunkedBodyWriter struct {
+	sync.Once
+	finalizeErr error
+	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.
+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 err = ext.WriteChunk(c.w, p, false); err != nil {
+		return
+	}
+	return len(p), nil
+}
+
+func (c *chunkedBodyWriter) Flush() error {
+	return c.w.Flush()
+}
+
+// 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() {
+		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
+}
+
+func NewChunkedBodyWriter(r *protocol.Response, w network.Writer) network.ExtWriter {
+	return &chunkedBodyWriter{
+		r: r,
+		w: w,
+	}
+}
diff --git a/pkg/protocol/http1/server.go b/pkg/protocol/http1/server.go
index 99d2da7..ba8915a 100644
--- a/pkg/protocol/http1/server.go
+++ b/pkg/protocol/http1/server.go
@@ -413,6 +413,11 @@ func writeErrorResponse(zw network.Writer, ctx *app.RequestContext, serverName [
 }
 
 func writeResponse(ctx *app.RequestContext, w network.Writer) error {
+	// Skip default response writing logic if it has been hijacked
+	if ctx.Response.GetHijackWriter() != nil {
+		return ctx.Response.GetHijackWriter().Finalize()
+	}
+
 	err := resp.Write(&ctx.Response, w)
 	if err != nil {
 		return err
diff --git a/pkg/protocol/response.go b/pkg/protocol/response.go
index 04440f3..016af92 100644
--- a/pkg/protocol/response.go
+++ b/pkg/protocol/response.go
@@ -46,6 +46,7 @@ import (
 	"net"
 	"sync"
 
+	"github.com/cloudwego/hertz/internal/bytesconv"
 	"github.com/cloudwego/hertz/internal/nocopy"
 	"github.com/cloudwego/hertz/pkg/common/bytebufferpool"
 	"github.com/cloudwego/hertz/pkg/common/compress"
@@ -90,6 +91,17 @@ type Response struct {
 	raddr net.Addr
 	// Local TCPAddr from concurrently net.Conn
 	laddr net.Addr
+
+	// If set a hijackWriter, hertz will skip the default header/body writer process.
+	hijackWriter network.ExtWriter
+}
+
+func (resp *Response) GetHijackWriter() network.ExtWriter {
+	return resp.hijackWriter
+}
+
+func (resp *Response) HijackWriter(writer network.ExtWriter) {
+	resp.hijackWriter = writer
 }
 
 type responseBodyWriter struct {
@@ -133,7 +145,7 @@ func (resp *Response) ConstructBodyStream(body *bytebufferpool.ByteBuffer, bodyS
 // BodyWriter returns writer for populating response body.
 //
 // If used inside RequestHandler, the returned writer must not be used
-// after returning from RequestHandler. Use RequestCtx.Write
+// after returning from RequestHandler. Use RequestContext.Write
 // or SetBodyStreamWriter in this case.
 func (resp *Response) BodyWriter() io.Writer {
 	resp.w.r = resp
@@ -265,6 +277,7 @@ func (resp *Response) Reset() {
 	resp.raddr = nil
 	resp.laddr = nil
 	resp.ImmediateHeaderFlush = false
+	resp.hijackWriter = nil
 }
 
 func (resp *Response) resetSkipHeader() {
@@ -302,8 +315,22 @@ func (resp *Response) StatusCode() int {
 //
 // It is safe re-using body argument after the function returns.
 func (resp *Response) SetBody(body []byte) {
-	resp.CloseBodyStream()      //nolint:errcheck
-	resp.BodyBuffer().Set(body) //nolint:errcheck
+	resp.CloseBodyStream() //nolint:errcheck
+	if resp.GetHijackWriter() == nil {
+		resp.BodyBuffer().Set(body) //nolint:errcheck
+		return
+	}
+
+	// If the hijack writer support .SetBody() api, then use it.
+	if setter, ok := resp.GetHijackWriter().(interface {
+		SetBody(b []byte)
+	}); ok {
+		setter.SetBody(body)
+		return
+	}
+
+	// Otherwise, call .Write() api instead.
+	resp.GetHijackWriter().Write(body) //nolint:errcheck
 }
 
 func (resp *Response) BodyStream() io.Reader {
@@ -314,13 +341,21 @@ func (resp *Response) BodyStream() io.Reader {
 //
 // It is safe re-using p after the function returns.
 func (resp *Response) AppendBody(p []byte) {
-	resp.CloseBodyStream()     //nolint:errcheck
+	resp.CloseBodyStream() //nolint:errcheck
+	if resp.hijackWriter != nil {
+		resp.hijackWriter.Write(p) //nolint:errcheck
+		return
+	}
 	resp.BodyBuffer().Write(p) //nolint:errcheck
 }
 
 // AppendBodyString appends s to response body.
 func (resp *Response) AppendBodyString(s string) {
-	resp.CloseBodyStream()           //nolint:errcheck
+	resp.CloseBodyStream() //nolint:errcheck
+	if resp.hijackWriter != nil {
+		resp.hijackWriter.Write(bytesconv.S2b(s)) //nolint:errcheck
+		return
+	}
 	resp.BodyBuffer().WriteString(s) //nolint:errcheck
 }
 


Test patch:
diff --git a/pkg/common/test/mock/network.go b/pkg/common/test/mock/network.go
index bb08cb4..22c8266 100644
--- a/pkg/common/test/mock/network.go
+++ b/pkg/common/test/mock/network.go
@@ -32,6 +32,12 @@ type Conn struct {
 	readTimeout time.Duration
 	zr          network.Reader
 	zw          network.ReadWriter
+	wroteLen    int
+}
+
+type Recorder interface {
+	network.Reader
+	WroteLen() int
 }
 
 func (m *Conn) SetWriteTimeout(t time.Duration) error {
@@ -98,19 +104,31 @@ func (m *Conn) Len() int {
 }
 
 func (m *Conn) Malloc(n int) (buf []byte, err error) {
+	m.wroteLen += n
 	return m.zw.Malloc(n)
 }
 
 func (m *Conn) WriteBinary(b []byte) (n int, err error) {
-	return m.zw.WriteBinary(b)
+	n, err = m.zw.WriteBinary(b)
+	m.wroteLen += n
+	return n, err
 }
 
 func (m *Conn) Flush() error {
 	return m.zw.Flush()
 }
 
-func (m *Conn) WriterRecorder() network.Reader {
-	return m.zw
+func (m *Conn) WriterRecorder() Recorder {
+	return &recorder{c: m, Reader: m.zw}
+}
+
+type recorder struct {
+	c *Conn
+	network.Reader
+}
+
+func (r *recorder) WroteLen() int {
+	return r.c.wroteLen
 }
 
 func (m *SlowReadConn) Peek(i int) ([]byte, error) {
diff --git a/pkg/common/test/mock/writer.go b/pkg/common/test/mock/writer.go
new file mode 100644
index 0000000..b326242
--- /dev/null
+++ b/pkg/common/test/mock/writer.go
@@ -0,0 +1,47 @@
+/*
+ * Copyright 2023 CloudWeGo Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package mock
+
+import "bytes"
+
+type ExtWriter struct {
+	tmp     []byte
+	Buf     *bytes.Buffer
+	IsFinal *bool
+}
+
+func (m *ExtWriter) Write(p []byte) (n int, err error) {
+	m.tmp = p
+	return len(p), nil
+}
+
+func (m *ExtWriter) Flush() error {
+	_, err := m.Buf.Write(m.tmp)
+	return err
+}
+
+func (m *ExtWriter) Finalize() error {
+	if !*m.IsFinal {
+		*m.IsFinal = true
+	}
+	return nil
+}
+
+func (m *ExtWriter) SetBody(body []byte) {
+	m.Buf.Reset()
+	m.tmp = body
+}
diff --git a/pkg/protocol/http1/resp/writer_test.go b/pkg/protocol/http1/resp/writer_test.go
new file mode 100644
index 0000000..bca2d43
--- /dev/null
+++ b/pkg/protocol/http1/resp/writer_test.go
@@ -0,0 +1,53 @@
+/*
+ * Copyright 2023 CloudWeGo Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package resp
+
+import (
+	"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"
+)
+
+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)))
+}
+
+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)))
+}
diff --git a/pkg/protocol/http1/server_test.go b/pkg/protocol/http1/server_test.go
index 4a2a15a..dfa8b78 100644
--- a/pkg/protocol/http1/server_test.go
+++ b/pkg/protocol/http1/server_test.go
@@ -17,6 +17,7 @@
 package http1
 
 import (
+	"bytes"
 	"context"
 	"errors"
 	"sync"
@@ -31,6 +32,8 @@ import (
 	"github.com/cloudwego/hertz/pkg/common/tracer/stats"
 	"github.com/cloudwego/hertz/pkg/common/tracer/traceinfo"
 	"github.com/cloudwego/hertz/pkg/network"
+	"github.com/cloudwego/hertz/pkg/protocol"
+	"github.com/cloudwego/hertz/pkg/protocol/http1/resp"
 )
 
 var pool = &sync.Pool{New: func() interface{} {
@@ -189,9 +192,65 @@ func TestEventStack(t *testing.T) {
 	}
 }
 
+func TestDefaultWriter(t *testing.T) {
+	server := &Server{}
+	reqCtx := &app.RequestContext{}
+	server.Core = &mockCore{
+		ctxPool: &sync.Pool{New: func() interface{} {
+			return reqCtx
+		}},
+		mockHandler: func(c context.Context, ctx *app.RequestContext) {
+			ctx.Write([]byte("hello, hertz"))
+			ctx.Flush()
+		},
+	}
+	defaultConn := mock.NewConn("GET / HTTP/1.1\nHost: foobar.com\n\n")
+	err := server.Serve(context.TODO(), defaultConn)
+	assert.True(t, errors.Is(err, errs.ErrShortConnection))
+	defaultResponseResult := defaultConn.WriterRecorder()
+	assert.DeepEqual(t, 0, defaultResponseResult.Len()) // all data is flushed so the buffer length is 0
+	response := protocol.AcquireResponse()
+	resp.Read(response, defaultResponseResult)
+	assert.DeepEqual(t, "hello, hertz", string(response.Body()))
+}
+
+func TestHijackResponseWriter(t *testing.T) {
+	server := &Server{}
+	reqCtx := &app.RequestContext{}
+	buf := new(bytes.Buffer)
+	isFinal := false
+	server.Core = &mockCore{
+		ctxPool: &sync.Pool{New: func() interface{} {
+			return reqCtx
+		}},
+		mockHandler: func(c context.Context, ctx *app.RequestContext) {
+			// response before write will be dropped
+			ctx.Write([]byte("invalid data"))
+
+			ctx.Response.HijackWriter(&mock.ExtWriter{
+				Buf:     buf,
+				IsFinal: &isFinal,
+			})
+
+			ctx.Write([]byte("hello, hertz"))
+			ctx.Flush()
+		},
+	}
+	defaultConn := mock.NewConn("GET / HTTP/1.1\nHost: foobar.com\n\n")
+	err := server.Serve(context.TODO(), defaultConn)
+	assert.True(t, errors.Is(err, errs.ErrShortConnection))
+	defaultResponseResult := defaultConn.WriterRecorder()
+	response := protocol.AcquireResponse()
+	resp.Read(response, defaultResponseResult)
+	assert.DeepEqual(t, 0, len(response.Body()))
+	assert.DeepEqual(t, "hello, hertz", buf.String())
+	assert.True(t, isFinal)
+}
+
 type mockCore struct {
-	ctxPool    *sync.Pool
-	controller tracer.Controller
+	ctxPool     *sync.Pool
+	controller  tracer.Controller
+	mockHandler func(c context.Context, ctx *app.RequestContext)
 }
 
 func (m *mockCore) IsRunning() bool {
@@ -202,7 +261,11 @@ func (m *mockCore) GetCtxPool() *sync.Pool {
 	return m.ctxPool
 }
 
-func (m *mockCore) ServeHTTP(c context.Context, ctx *app.RequestContext) {}
+func (m *mockCore) ServeHTTP(c context.Context, ctx *app.RequestContext) {
+	if m.mockHandler != nil {
+		m.mockHandler(c, ctx)
+	}
+}
 
 func (m *mockCore) GetTracer() tracer.Controller {
 	return m.controller
diff --git a/pkg/protocol/response_test.go b/pkg/protocol/response_test.go
index 1ecb806..20a18ff 100644
--- a/pkg/protocol/response_test.go
+++ b/pkg/protocol/response_test.go
@@ -51,6 +51,7 @@ import (
 	"github.com/cloudwego/hertz/pkg/common/bytebufferpool"
 	"github.com/cloudwego/hertz/pkg/common/compress"
 	"github.com/cloudwego/hertz/pkg/common/test/assert"
+	"github.com/cloudwego/hertz/pkg/common/test/mock"
 	"github.com/cloudwego/hertz/pkg/protocol/consts"
 )
 
@@ -272,3 +273,24 @@ func TestRespSafeCopy(t *testing.T) {
 		assert.DeepEqual(t, []byte{byte(i)}, resps[i].Body())
 	}
 }
+
+func TestResponse_HijackWriter(t *testing.T) {
+	resp := AcquireResponse()
+	buf := new(bytes.Buffer)
+	isFinal := false
+	resp.HijackWriter(&mock.ExtWriter{Buf: buf, IsFinal: &isFinal})
+	resp.AppendBody([]byte("hello"))
+	assert.DeepEqual(t, 0, buf.Len())
+	resp.GetHijackWriter().Flush()
+	assert.DeepEqual(t, "hello", buf.String())
+	resp.AppendBodyString(", world")
+	assert.DeepEqual(t, "hello", buf.String())
+	resp.GetHijackWriter().Flush()
+	assert.DeepEqual(t, "hello, world", buf.String())
+	resp.SetBody([]byte("hello, hertz"))
+	resp.GetHijackWriter().Flush()
+	assert.DeepEqual(t, "hello, hertz", buf.String())
+	assert.False(t, isFinal)
+	resp.GetHijackWriter().Finalize()
+	assert.True(t, isFinal)
+}