diff --git a/ctx.go b/ctx.go index 389fec99..8e285f6e 100644 --- a/ctx.go +++ b/ctx.go @@ -1014,12 +1014,18 @@ func (c *Ctx) Vary(fields ...string) { c.Append(HeaderVary, fields...) } -// Write writes p into response body. -func (c *Ctx) Write(p []byte) (n int, err error) { +// Write appends p into response body. +func (c *Ctx) Write(p []byte) (int, error) { c.fasthttp.Response.AppendBody(p) return len(p), nil } +// WriteString appends s to response body. +func (c *Ctx) WriteString(s string) (int, error) { + c.fasthttp.Response.AppendBodyString(s) + return len(s), nil +} + // XHR returns a Boolean property, that is true, if the request's X-Requested-With header field is XMLHttpRequest, // indicating that the request was issued by a client library (such as jQuery). func (c *Ctx) XHR() bool { diff --git a/ctx_test.go b/ctx_test.go index 4bcd67c5..d4a97ce3 100644 --- a/ctx_test.go +++ b/ctx_test.go @@ -1936,3 +1936,52 @@ func Benchmark_Ctx_QueryParser(b *testing.B) { } utils.AssertEqual(b, nil, c.QueryParser(q)) } + +// go test -run Test_Ctx_BodyStreamWriter +func Test_Ctx_BodyStreamWriter(t *testing.T) { + t.Parallel() + + ctx := &fasthttp.RequestCtx{} + + ctx.SetBodyStreamWriter(func(w *bufio.Writer) { + fmt.Fprintf(w, "body writer line 1\n") + if err := w.Flush(); err != nil { + t.Errorf("unexpected error: %s", err) + } + fmt.Fprintf(w, "body writer line 2\n") + }) + if !ctx.IsBodyStream() { + t.Fatal("IsBodyStream must return true") + } + + s := ctx.Response.String() + br := bufio.NewReader(bytes.NewBufferString(s)) + var resp fasthttp.Response + if err := resp.Read(br); err != nil { + t.Fatalf("Error when reading response: %s", err) + } + body := string(resp.Body()) + expectedBody := "body writer line 1\nbody writer line 2\n" + if body != expectedBody { + t.Fatalf("unexpected body: %q. Expecting %q", body, expectedBody) + } +} + +// go test -v -run=^$ -bench=Benchmark_Ctx_BodyStreamWriter -benchmem -count=4 +func Benchmark_Ctx_BodyStreamWriter(b *testing.B) { + ctx := &fasthttp.RequestCtx{} + user := []byte(`{"name":"john"}`) + b.ReportAllocs() + b.ResetTimer() + for n := 0; n < b.N; n++ { + ctx.ResetBody() + ctx.SetBodyStreamWriter(func(w *bufio.Writer) { + for i := 0; i < 10; i++ { + w.Write(user) + if err := w.Flush(); err != nil { + return + } + } + }) + } +}