Skip to content

Commit

Permalink
Add optimized ReadAll for http responses
Browse files Browse the repository at this point in the history
1. When the body content length is known it's much faster to
pre-allocated the entire buffer once.
2. When the length is unknown using `bytes.NewBuffer` + `io.Copy` is
much faster.

Also, added benchmarks to prove the difference.
  • Loading branch information
rdner committed Feb 12, 2024
1 parent 4739839 commit 212ef1b
Show file tree
Hide file tree
Showing 2 changed files with 153 additions and 0 deletions.
38 changes: 38 additions & 0 deletions transport/httpcommon/httpcommon.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,10 @@
package httpcommon

import (
"bytes"
"errors"
"fmt"
"io"
"net/http"
"time"

Expand Down Expand Up @@ -410,3 +414,37 @@ func WithLogger(logger *logp.Logger) TransportOption {
s.logger = logger
})
}

// ReadAll returns the whole response body as bytes.
// This is an optimized version of `io.ReadAll`.
func ReadAll(resp *http.Response) ([]byte, error) {
if resp == nil {
return nil, errors.New("response cannot be nil")
}
switch {
case resp.ContentLength == 0:
return []byte{}, nil
// if we know the body length we can allocate the buffer only once
case resp.ContentLength >= 0:
body := make([]byte, resp.ContentLength)
_, err := io.ReadFull(resp.Body, body)
if err != nil {
return nil, fmt.Errorf("failed to read the response body with a known length %d: %w", resp.ContentLength, err)
}
return body, nil

default:
// using `bytes.NewBuffer` + `io.Copy` is much faster than `io.ReadAll`
// see https://github.com/elastic/beats/issues/36151#issuecomment-1931696767
buf := bytes.NewBuffer(nil)
_, err := io.Copy(buf, resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read the response body with unknown length: %w", err)
}
body := buf.Bytes()
if body == nil {
body = []byte{}
}
return body, nil
}
}
115 changes: 115 additions & 0 deletions transport/httpcommon/httpcommon_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,10 @@
package httpcommon

import (
"bytes"
"fmt"
"io"
"net/http"
"testing"
"time"

Expand Down Expand Up @@ -92,3 +96,114 @@ ssl:
})
}
}

func TestReadAll(t *testing.T) {
size := 100
body := bytes.Repeat([]byte{'a'}, size)
cases := []struct {
name string
resp *http.Response
expBody []byte
}{
{
name: "reads known size",
resp: &http.Response{
ContentLength: int64(size),
Body: io.NopCloser(bytes.NewBuffer(body)),
},
expBody: body,
},
{
name: "reads unknown size",
resp: &http.Response{
ContentLength: -1,
Body: io.NopCloser(bytes.NewBuffer(body)),
},
expBody: body,
},
{
name: "supports empty with size=0",
resp: &http.Response{
ContentLength: 0,
Body: io.NopCloser(bytes.NewBuffer(nil)),
},
expBody: []byte{},
},
{
name: "supports empty with unknown size",
resp: &http.Response{
ContentLength: -1,
Body: io.NopCloser(bytes.NewBuffer(nil)),
},
expBody: []byte{},
},
}

for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
actBody, err := ReadAll(tc.resp)
require.NoError(t, err)
require.Equal(t, tc.expBody, actBody)
})
}
}

func BenchmarkReadAll(b *testing.B) {
sizes := []int{
100, // 100 bytes
100 * 1024, // 100KB
1024 * 1024, // 1MB
}
for _, size := range sizes {
b.Run(fmt.Sprintf("size: %d", size), func(b *testing.B) {

// emulate a file or an HTTP response
generated := bytes.Repeat([]byte{'a'}, size)
content := bytes.NewReader(generated)
cases := []struct {
name string
resp *http.Response
}{
{
name: "unknown length",
resp: &http.Response{
ContentLength: -1,
Body: io.NopCloser(content),
},
},
{
name: "known length",
resp: &http.Response{
ContentLength: int64(size),
Body: io.NopCloser(content),
},
},
}

b.ResetTimer()

for _, tc := range cases {
b.Run(tc.name, func(b *testing.B) {
b.Run("io.ReadAll", func(b *testing.B) {
for i := 0; i < b.N; i++ {
_, err := content.Seek(0, io.SeekStart) // reset
require.NoError(b, err)
data, err := io.ReadAll(tc.resp.Body)
require.NoError(b, err)
require.Equalf(b, size, len(data), "size does not match, expected %d, actual %d", size, len(data))
}
})
b.Run("bytes.Buffer+io.Copy", func(b *testing.B) {
for i := 0; i < b.N; i++ {
_, err := content.Seek(0, io.SeekStart) // reset
require.NoError(b, err)
data, err := ReadAll(tc.resp)
require.NoError(b, err)
require.Equalf(b, size, len(data), "size does not match, expected %d, actual %d", size, len(data))
}
})
})
}
})
}
}

0 comments on commit 212ef1b

Please sign in to comment.