Add optimized ReadAll function

api/npipe/listener_windows_test.go

@@ -21,10 +21,10 @@ package npipe
 
 import (
  "fmt"
- "io/ioutil"
  "net/http"
  "testing"
 
+ "github.com/elastic/elastic-agent-libs/transport/httpcommon"
  "github.com/stretchr/testify/assert"
  "github.com/stretchr/testify/require"
 )
@@ -43,7 +43,8 @@ func TestHTTPOverNamedPipe(t *testing.T) {
  })
 
  go func() {
- _ = http.Serve(l, mux)
+ _ = http.Serve(l, mux) //nolint:gosec // Serve does not support setting timeouts, it is fine for tests.
+
  }()
 
  c := http.Client{
@@ -52,10 +53,10 @@ func TestHTTPOverNamedPipe(t *testing.T) {
  },
  }
 
- // nolint:noctx // for testing purposes
+ //nolint:noctx // for testing purposes
  r, err := c.Get("http://npipe/echo-hello")
  require.NoError(t, err)
- body, err := ioutil.ReadAll(r.Body)
+ body, err := httpcommon.ReadAll(r)
  require.NoError(t, err)
  defer r.Body.Close()
 

api/server_test.go

@@ -20,8 +20,6 @@ package api
 import (
  "context"
  "fmt"
- "io"
- "io/ioutil"
  "net"
  "net/http"
  "os"
@@ -32,6 +30,7 @@ import (
  "github.com/stretchr/testify/require"
 
  "github.com/elastic/elastic-agent-libs/config"
+ "github.com/elastic/elastic-agent-libs/transport/httpcommon"
 )
 
 const (
@@ -71,10 +70,7 @@ func TestSocket(t *testing.T) {
  }
 
  t.Run("socket doesn't exist before", func(t *testing.T) {
- tmpDir, err := ioutil.TempDir("", "testsocket")
- require.NoError(t, err)
- defer os.RemoveAll(tmpDir)
-
+ tmpDir := t.TempDir()
  sockFile := tmpDir + "/test.sock"
 
  cfg := config.MustNewConfigFrom(map[string]interface{}{
@@ -102,11 +98,7 @@ func TestSocket(t *testing.T) {
  })
 
  t.Run("starting beat and recover a dangling socket file", func(t *testing.T) {
- tmpDir, err := ioutil.TempDir("", "testsocket")
- require.NoError(t, err)
- defer os.RemoveAll(tmpDir)
-
- sockFile := tmpDir + "/test.sock"
+ sockFile := t.TempDir() + "/test.sock"
 
  // Create the socket before the server.
  f, err := os.Create(sockFile)
@@ -161,7 +153,7 @@ func getResponse(t *testing.T, sockFile, url string) string {
  require.NoError(t, err)
  defer r.Body.Close()
 
- body, err := ioutil.ReadAll(r.Body)
+ body, err := httpcommon.ReadAll(r)
  require.NoError(t, err)
  return string(body)
 }
@@ -188,7 +180,7 @@ func TestHTTP(t *testing.T) {
  require.NoError(t, err)
  }()
 
- body, err := ioutil.ReadAll(r.Body)
+ body, err := httpcommon.ReadAll(r)
  require.NoError(t, err)
 
  assert.Equal(t, "ehlo!", string(body))
@@ -229,7 +221,7 @@ func TestAttachHandler(t *testing.T) {
  require.NoError(t, err)
  }()
 
- body, err := io.ReadAll(r.Body)
+ body, err := httpcommon.ReadAll(r)
  require.NoError(t, err)
 
  assert.Equal(t, "test!", string(body))

api/server_windows_test.go

@@ -20,7 +20,6 @@
 package api
 
 import (
- "io/ioutil"
  "net/http"
  "testing"
 
@@ -29,6 +28,7 @@ import (
 
  "github.com/elastic/elastic-agent-libs/api/npipe"
  "github.com/elastic/elastic-agent-libs/config"
+ "github.com/elastic/elastic-agent-libs/transport/httpcommon"
 )
 
 func TestNamedPipe(t *testing.T) {
@@ -52,12 +52,12 @@ func TestNamedPipe(t *testing.T) {
  },
  }
 
- // nolint:noctx // for testing purposes
+ //nolint:noctx // for testing purposes
  r, err := c.Get("http://npipe/echo-hello")
  require.NoError(t, err)
  defer r.Body.Close()
 
- body, err := ioutil.ReadAll(r.Body)
+ body, err := httpcommon.ReadAll(r)
  require.NoError(t, err)
 
  assert.Equal(t, "ehlo!", string(body))

file/fileinfo_test.go

@@ -24,7 +24,6 @@
 package file_test
 
 import (
- "io/ioutil"
  "os"
  "path/filepath"
  "testing"
@@ -35,7 +34,7 @@ import (
 )
 
 func TestStat(t *testing.T) {
- f, err := ioutil.TempFile("", "teststat")
+ f, err := os.CreateTemp("", "teststat")
  if err != nil {
  t.Fatal(err)
  }

file/helper_test.go

@@ -20,7 +20,7 @@
 package file
 
 import (
- "io/ioutil"
+ "fmt"
  "os"
  "path/filepath"
  "testing"
@@ -29,19 +29,15 @@ import (
 )
 
 func TestSafeFileRotateExistingFile(t *testing.T) {
- tempdir, err := ioutil.TempDir("", "")
- assert.NoError(t, err)
- defer func() {
- assert.NoError(t, os.RemoveAll(tempdir))
- }()
+ tempdir := t.TempDir()
 
  // create an existing registry file
- err = ioutil.WriteFile(filepath.Join(tempdir, "registry"),
+ err := os.WriteFile(filepath.Join(tempdir, "registry"),
  []byte("existing filebeat"), 0x777)
  assert.NoError(t, err)
 
  // create a new registry.new file
- err = ioutil.WriteFile(filepath.Join(tempdir, "registry.new"),
+ err = os.WriteFile(filepath.Join(tempdir, "registry.new"),
  []byte("new filebeat"), 0x777)
  assert.NoError(t, err)
 
@@ -50,35 +46,23 @@ func TestSafeFileRotateExistingFile(t *testing.T) {
  filepath.Join(tempdir, "registry.new"))
  assert.NoError(t, err)
 
- contents, err := ioutil.ReadFile(filepath.Join(tempdir, "registry"))
+ contents, err := os.ReadFile(filepath.Join(tempdir, "registry"))
  assert.NoError(t, err)
  assert.Equal(t, []byte("new filebeat"), contents)
 
- // do it again to make sure we deal with deleting the old file
-
- err = ioutil.WriteFile(filepath.Join(tempdir, "registry.new"),
- []byte("new filebeat 1"), 0x777)
- assert.NoError(t, err)
-
- err = SafeFileRotate(filepath.Join(tempdir, "registry"),
- filepath.Join(tempdir, "registry.new"))
- assert.NoError(t, err)
-
- contents, err = ioutil.ReadFile(filepath.Join(tempdir, "registry"))
- assert.NoError(t, err)
- assert.Equal(t, []byte("new filebeat 1"), contents)
-
- // and again for good measure
-
- err = ioutil.WriteFile(filepath.Join(tempdir, "registry.new"),
- []byte("new filebeat 2"), 0x777)
- assert.NoError(t, err)
-
- err = SafeFileRotate(filepath.Join(tempdir, "registry"),
- filepath.Join(tempdir, "registry.new"))
- assert.NoError(t, err)
-
- contents, err = ioutil.ReadFile(filepath.Join(tempdir, "registry"))
- assert.NoError(t, err)
- assert.Equal(t, []byte("new filebeat 2"), contents)
+ // do it twice to make sure we deal with deleting the old file
+ for i := 0; i < 2; i++ {
+ expectedContents := []byte(fmt.Sprintf("new filebeat %d", i))
+ err = os.WriteFile(filepath.Join(tempdir, "registry.new"),
+ expectedContents, 0x777)
+ assert.NoError(t, err)
+
+ err = SafeFileRotate(filepath.Join(tempdir, "registry"),
+ filepath.Join(tempdir, "registry.new"))
+ assert.NoError(t, err)
+
+ contents, err = os.ReadFile(filepath.Join(tempdir, "registry"))
+ assert.NoError(t, err)
+ assert.Equal(t, expectedContents, contents)
+ }
 }

iobuf/iobuf.go

@@ -0,0 +1,36 @@
+// Licensed to Elasticsearch B.V. under one or more contributor
+// license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright
+// ownership. Elasticsearch B.V. licenses this file to you under
+// the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package iobuf
+
+import (
+ "bytes"
+ "io"
+)
+
+// ReadAll reads all data from r and returns it as a byte slice.
+// A successful call returns err == nil, not err == EOF. It does not
+// treat an EOF as an error to be reported.
+//
+// This function is similar to io.ReadAll, but uses a bytes.Buffer to
+// accumulate the data, which has a more efficient growing algorithm and
+// uses io.WriterTo if r implements it.
+func ReadAll(r io.Reader) ([]byte, error) {
+ var buf bytes.Buffer
+ _, err := io.Copy(&buf, r)
+ return buf.Bytes(), err
+}

iobuf/iobuf_test.go

@@ -0,0 +1,123 @@
+// Licensed to Elasticsearch B.V. under one or more contributor
+// license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright
+// ownership. Elasticsearch B.V. licenses this file to you under
+// the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package iobuf
+
+import (
+ "bytes"
+ "fmt"
+ "io"
+ "log"
+ "strings"
+ "testing"
+)
+
+func ExampleReadAll() {
+ r := strings.NewReader("The quick brown fox jumps over the lazy dog.")
+
+ b, err := ReadAll(r)
+ if err != nil {
+ log.Fatal(err)
+ }
+
+ fmt.Printf("%s", b)
+
+ // Output:
+ // The quick brown fox jumps over the lazy dog.
+}
+
+// dumbReadSeeker is a ReadSeeker that does not implement the io.WriteTo optimization.
+type dumbReadSeeker struct {
+ rs io.ReadSeeker
+}
+
+func (r *dumbReadSeeker) Read(p []byte) (n int, err error) {
+ return r.rs.Read(p)
+}
+
+func (r *dumbReadSeeker) Seek(offset int64, whence int) (int64, error) {
+ return r.rs.Seek(offset, whence)
+}
+
+func genData(n int) io.ReadSeeker {
+ return bytes.NewReader(bytes.Repeat([]byte{'a'}, n))
+}
+
+func genDataDumb(n int) io.ReadSeeker {
+ return &dumbReadSeeker{rs: genData(n)}
+}
+
+func BenchmarkReadAll(b *testing.B) {
+ // Make sure we test different sizes to overcome initial buffer sizes:
+ // io.ReadAll uses a 512 bytes buffer
+ // bytes.Buffer uses a 64 bytes buffer
+ sizes := []int{
+ 32, // 32 bytes
+ 64, // 64 bytes
+ 512, // 512 bytes
+ 10 * 1024, // 10KB
+ 100 * 1024, // 100KB
+ 1024 * 1024, // 1MB
+ }
+ sizesReadable := []string{
+ "32B",
+ "64B",
+ "512B",
+ "10KB",
+ "100KB",
+ "1MB",
+ }
+
+ benchFunc := func(b *testing.B, size int, genFunc func(n int) io.ReadSeeker, readFunc func(io.Reader) ([]byte, error)) {
+ buf := genFunc(size)
+ b.ResetTimer()
+ for i := 0; i < b.N; i++ {
+ _, err := buf.Seek(0, io.SeekStart) // reset
+ if err != nil {
+ b.Fatal(err)
+ }
+ data, err := readFunc(buf)
+ if err != nil {
+ b.Fatal(err)
+ }
+ if len(data) != size {
+ b.Fatalf("size does not match, expected %d, actual %d", size, len(data))
+ }
+ }
+ }
+
+ for i, size := range sizes {
+ b.Run(fmt.Sprintf("size %s", sizesReadable[i]), func(b *testing.B) {
+ b.Run("io.ReadAll", func(b *testing.B) {
+ b.Run("WriterTo", func(b *testing.B) {
+ benchFunc(b, size, genData, io.ReadAll)
+ })
+ b.Run("Reader", func(b *testing.B) {
+ benchFunc(b, size, genDataDumb, io.ReadAll)
+ })
+ })
+ b.Run("ReadAll", func(b *testing.B) {
+ b.Run("WriterTo", func(b *testing.B) {
+ benchFunc(b, size, genData, ReadAll)
+ })
+ b.Run("Reader", func(b *testing.B) {
+ benchFunc(b, size, genDataDumb, ReadAll)
+ })
+ })
+ })
+ }
+}