Repository: valyala/fasthttp
Pull request #477 (2018-11-19 12:25:14)

Support huge read-only []byte response bodies
See issue #469.


erikdubbelboer: @cipriancraciun would it be possible for you to add some tests to make sure everything keeps working correctly in the future?
cipriancraciun: > @cipriancraciun would it be possible for you to add some tests to make sure everything keeps working correctly in the future?

Sorry for deferring this, but I was a little busy.  I'll try to add them during these days around new year's eve.

cipriancraciun: So I've rebased my two commits based on the latest code, and added a few tests at the end of `http_tests.go`.  Basically I check that the `SetBodyRaw` actually works as expected, if `CopyTo` and `ResetBody` work properly.

If I can make any other changes please let me know.


Patch:
diff --git a/http.go b/http.go
index 10dc465..79241db 100644
--- a/http.go
+++ b/http.go
@@ -67,6 +67,7 @@ type Response struct {
 	bodyStream io.Reader
 	w          responseBodyWriter
 	body       *bytebufferpool.ByteBuffer
+	bodyRaw    []byte
 
 	// Response.Read() skips reading body if set to true.
 	// Use it for reading HEAD responses.
@@ -321,6 +322,9 @@ func (resp *Response) Body() []byte {
 }
 
 func (resp *Response) bodyBytes() []byte {
+	if resp.bodyRaw != nil {
+		return resp.bodyRaw
+	}
 	if resp.body == nil {
 		return nil
 	}
@@ -338,6 +342,7 @@ func (resp *Response) bodyBuffer() *bytebufferpool.ByteBuffer {
 	if resp.body == nil {
 		resp.body = responseBodyPool.Get()
 	}
+	resp.bodyRaw = nil
 	return resp.body
 }
 
@@ -462,6 +467,7 @@ func (resp *Response) SetBodyString(body string) {
 
 // ResetBody resets response body.
 func (resp *Response) ResetBody() {
+	resp.bodyRaw = nil
 	resp.closeBodyStream()
 	if resp.body != nil {
 		if resp.keepBodyBuffer {
@@ -473,6 +479,14 @@ func (resp *Response) ResetBody() {
 	}
 }
 
+// SetBodyRaw sets response body, but without copying it.
+//
+// From this point onward the body argument must not be changed.
+func (resp *Response) SetBodyRaw(body []byte) {
+	resp.ResetBody()
+	resp.bodyRaw = body
+}
+
 // ReleaseBody retires the response body if it is greater than "size" bytes.
 //
 // This permits GC to reclaim the large buffer.  If used, must be before
@@ -481,6 +495,7 @@ func (resp *Response) ResetBody() {
 // Use this method only if you really understand how it works.
 // The majority of workloads don't need this method.
 func (resp *Response) ReleaseBody(size int) {
+	resp.bodyRaw = nil
 	if cap(resp.body.B) > size {
 		resp.closeBodyStream()
 		resp.body = nil
@@ -519,6 +534,8 @@ func (resp *Response) SwapBody(body []byte) []byte {
 		}
 	}
 
+	resp.bodyRaw = nil
+
 	oldBody := bb.B
 	bb.B = body
 	return oldBody
@@ -639,7 +656,12 @@ func (req *Request) copyToSkipBody(dst *Request) {
 // CopyTo copies resp contents to dst except of body stream.
 func (resp *Response) CopyTo(dst *Response) {
 	resp.copyToSkipBody(dst)
-	if resp.body != nil {
+	if resp.bodyRaw != nil {
+		dst.bodyRaw = resp.bodyRaw
+		if dst.body != nil {
+			dst.body.Reset()
+		}
+	} else if resp.body != nil {
 		dst.bodyBuffer().Set(resp.body.B)
 	} else if dst.body != nil {
 		dst.body.Reset()
@@ -661,6 +683,7 @@ func swapRequestBody(a, b *Request) {
 
 func swapResponseBody(a, b *Response) {
 	a.body, b.body = b.body, a.body
+	a.bodyRaw, b.bodyRaw = b.bodyRaw, a.bodyRaw
 	a.bodyStream, b.bodyStream = b.bodyStream, a.bodyStream
 }
 
@@ -1265,6 +1288,7 @@ func (resp *Response) gzipBody(level int) error {
 			responseBodyPool.Put(resp.body)
 		}
 		resp.body = w
+		resp.bodyRaw = nil
 	}
 	resp.Header.SetCanonical(strContentEncoding, strGzip)
 	return nil
@@ -1319,6 +1343,7 @@ func (resp *Response) deflateBody(level int) error {
 			responseBodyPool.Put(resp.body)
 		}
 		resp.body = w
+		resp.bodyRaw = nil
 	}
 	resp.Header.SetCanonical(strContentEncoding, strDeflate)
 	return nil


Test patch:
diff --git a/http_test.go b/http_test.go
index 3253e95..653142c 100644
--- a/http_test.go
+++ b/http_test.go
@@ -1891,3 +1891,35 @@ Content-Type: application/json
 		t.Fatalf("unexpected output %q", w.Bytes())
 	}
 }
+
+
+func TestResponseRawBodySet(t *testing.T) {
+	var resp Response
+	
+	expectedS := "test"
+	body := []byte(expectedS)
+	resp.SetBodyRaw(body)
+	
+	testBodyWriteTo(t, &resp, expectedS, true)
+}
+
+func TestResponseRawBodyReset(t *testing.T) {
+	var resp Response
+	
+	body := []byte("test")
+	resp.SetBodyRaw(body)
+	resp.ResetBody()
+	
+	testBodyWriteTo(t, &resp, "", true)
+}
+
+func TestResponseRawBodyCopyTo(t *testing.T) {
+	var resp Response
+	
+	expectedS := "test"
+	body := []byte(expectedS)
+	resp.SetBodyRaw(body)
+	
+	testResponseCopyTo(t, &resp)
+}
+