Skip to content

Commit

Permalink
Add optimized ReadAll for http responses (#183)
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 13, 2024
1 parent d9ce6d3 commit a275281
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 a275281

Please sign in to comment.