Improve memory consumption for length calculation (#564) (#565)

Add tests
Update tests that relies on a specific set of request

Co-authored-by: Laurent Saint-Félix <[email protected]>

Author: github-actions[bot]

esutil/bulk_indexer.go

@@ -21,6 +21,7 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -179,22 +180,25 @@ func (item *BulkIndexerItem) marshallMeta() {
 // computeLength calculate the size of the body and the metadata.
 func (item *BulkIndexerItem) computeLength() error {
 	if item.Body != nil {
-		var buf bytes.Buffer
-		_, err := io.Copy(&buf, item.Body)
-		if err != nil {
-			return err
+		// TODO propagate buf len to config to allow for performance gains.
+		var buf = make([]byte, 1<<4)
+		for {
+			n, err := item.Body.Read(buf)
+			if errors.Is(err, io.EOF) {
+				break
+			} else if err != nil {
+				return err
+			}
+			item.payloadLength += n
 		}
-
-		_, err = item.Body.Seek(0, io.SeekStart)
+		_, err := item.Body.Seek(0, io.SeekStart)
 		if err != nil {
 			return err
 		}
-		item.payloadLength = buf.Len()
-		return nil
 	}
 	item.payloadLength += len(item.meta.Bytes())
-	// Add two bytes to account for newlines.
-	item.payloadLength += 2
+	// Add one byte to account for newline at the end of payload.
+	item.payloadLength++
 
 	return nil
 }

esutil/bulk_indexer_internal_test.go

@@ -99,7 +99,7 @@ func TestBulkIndexer(t *testing.T) {
 
 		cfg := BulkIndexerConfig{
 			NumWorkers:    1,
-			FlushBytes:    38 * 2,    // 38 bytes header + body, times 2 to match 2 responses per file in testdata
+			FlushBytes:    39 * 2,    // 38 bytes header + body, times 2 to match 2 responses per file in testdata
 			FlushInterval: time.Hour, // Disable auto-flushing, because response doesn't match number of items
 			Client:        es}
 		if os.Getenv("DEBUG") != "" {
@@ -576,7 +576,7 @@ func TestBulkIndexer(t *testing.T) {
 		}
 		es, _ := elasticsearch.NewClient(esCfg)
 
-		biCfg := BulkIndexerConfig{NumWorkers: 1, FlushBytes: 50, Client: es}
+		biCfg := BulkIndexerConfig{NumWorkers: 1, FlushBytes: 28*2, Client: es}
 		if os.Getenv("DEBUG") != "" {
 			biCfg.DebugLogger = log.New(os.Stdout, "", 0)
 		}
@@ -955,6 +955,36 @@ func TestBulkIndexer(t *testing.T) {
 	})
 }
 
+func TestBulkIndexerItem(t *testing.T) {
+	body := `{"body":"Lorem ipsum dolor sit amet, consectetur adipisicing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat."}`
+	t.Run("correct computeLength size", func(t *testing.T) {
+		expectedLength := 266
+		bi := BulkIndexerItem{
+			Action:     "index",
+			DocumentID: strconv.Itoa(1),
+			Body:       strings.NewReader(body),
+		}
+		bi.marshallMeta()
+		bi.computeLength()
+		if bi.payloadLength != expectedLength {
+			t.Fatalf("invalid length, expected %d, got %d", expectedLength, bi.payloadLength)
+		}
+	})
+	t.Run("empty reader length should be meta length plus newlines", func(t *testing.T) {
+		expectedLength := 23
+		bi := BulkIndexerItem{
+			Action:     "index",
+			DocumentID: strconv.Itoa(1),
+			Body:       strings.NewReader(""),
+		}
+		bi.marshallMeta()
+		bi.computeLength()
+		if bi.payloadLength != expectedLength {
+			t.Fatalf("invalid length, expected %d, got %d", expectedLength, bi.payloadLength)
+		}
+	})
+}
+
 type customJSONDecoder struct{}
 
 func (d customJSONDecoder) UnmarshalFromReader(r io.Reader, blk *BulkIndexerResponse) error {