Support using io.Reader for bodies

client.go

@@ -1137,17 +1137,32 @@ func (c *HostClient) doNonNilReqResp(req *Request, resp *Response) (bool, error)
 	}
 
 	br := c.acquireReader(conn)
-	if err = resp.ReadLimitBody(br, c.MaxResponseBodySize); err != nil {
+	if resp.ReturnBodyReader {
+		if resp.BodyReader, err = resp.ReadLimitBodyReader(br, c.MaxResponseBodySize); err != nil {
+			c.releaseReader(br)
+			c.closeConn(cc)
+			return true, err
+		}
+		wrap := wrapBodyReaderPool.Get().(*wrapBodyReader)
+		wrap.r = resp.BodyReader
+		wrap.br = br
+		wrap.c = c
+		wrap.cc = cc
+		wrap.close = resetConnection || req.ConnectionClose() || resp.ConnectionClose()
+		resp.BodyReader = wrap
+	} else {
+		if err = resp.ReadLimitBody(br, c.MaxResponseBodySize); err != nil {
+			c.releaseReader(br)
+			c.closeConn(cc)
+			return true, err
+		}
 		c.releaseReader(br)
-		c.closeConn(cc)
-		return true, err
-	}
-	c.releaseReader(br)
 
-	if resetConnection || req.ConnectionClose() || resp.ConnectionClose() {
-		c.closeConn(cc)
-	} else {
-		c.releaseConn(cc)
+		if resetConnection || req.ConnectionClose() || resp.ConnectionClose() {
+			c.closeConn(cc)
+		} else {
+			c.releaseConn(cc)
+		}
 	}
 
 	return false, err

erik_test.go

@@ -0,0 +1,99 @@
+package fasthttp
+
+import (
+	"bytes"
+	"net"
+	//"os"
+	//"runtime"
+	"testing"
+
+	//"github.com/erikdubbelboer/bpprof"
+	"github.com/valyala/fasthttp/fasthttputil"
+)
+
+func BenchmarkErikHttpSimple(b *testing.B) {
+	//runtime.MemProfileRate = 1
+
+	body := []byte("test")
+	ln := fasthttputil.NewInmemoryListener()
+
+	go Serve(ln, func(ctx *RequestCtx) {
+		ctx.SetStatusCode(200)
+		ctx.SetBody(body)
+	})
+
+	c := Client{
+		Dial: func(addr string) (net.Conn, error) {
+			return ln.Dial()
+		},
+	}
+
+	var reuseBody []byte
+	for i := 0; i < b.N; i++ {
+		_, responseBody, err := c.Get(reuseBody, "http://example.com")
+		if err != nil {
+			panic(err)
+		}
+		reuseBody = responseBody
+	}
+
+	//bpprof.Heap(os.Stderr, "allocobjects")
+}
+
+func BenchmarkErikHttpBodyReader(b *testing.B) {
+	//runtime.MemProfileRate = 1
+
+	requestBody := []byte("request")
+	responseBody := []byte("response")
+	ln := fasthttputil.NewInmemoryListener()
+
+	go Serve(ln, func(ctx *RequestCtx) {
+		if !bytes.Equal(requestBody, ctx.PostBody()) {
+			panic("request body is wrong")
+		}
+		ctx.SetStatusCode(200)
+		ctx.SetBody(responseBody)
+	})
+
+	c := Client{
+		Dial: func(addr string) (net.Conn, error) {
+			return ln.Dial()
+		},
+	}
+
+	var body bytes.Buffer
+
+	requestURI := []byte("http://example.com")
+	requestMethod := []byte("POST")
+
+	for i := 0; i < b.N; i++ {
+		req := AcquireRequest()
+		res := AcquireResponse()
+
+		req.SetRequestURIBytes(requestURI)
+		req.SetBody(requestBody)
+		req.Header.SetMethodBytes(requestMethod)
+
+		res.ReturnBodyReader = true
+
+		if err := c.Do(req, res); err != nil {
+			panic(err)
+		}
+
+		body.Reset()
+		if _, err := body.ReadFrom(res.BodyReader); err != nil {
+			panic(err)
+		}
+
+		if !bytes.Equal(responseBody, body.Bytes()) {
+			panic("response body is wrong")
+		}
+
+		res.BodyReader.Close()
+		ReleaseBodyReader(res.BodyReader)
+		ReleaseRequest(req)
+		ReleaseResponse(res)
+	}
+
+	//bpprof.Heap(os.Stderr, "allocobjects")
+}