[BulkIndexer] reduce memory consumption (#558) (#559)

* Improve BulkIndexer memory consumption

* Flush instantly oversized payloads

Co-authored-by: Laurent Saint-Félix <laurent.saintfelix@elastic.co>

Author: github-actions[bot]

esutil/bulk_indexer.go

@@ -37,7 +37,6 @@ import (
 )
 
 // BulkIndexer represents a parallel, asynchronous, efficient indexer for Elasticsearch.
-//
 type BulkIndexer interface {
 	// Add adds an item to the indexer. It returns an error when the item cannot be added.
 	// Use the OnSuccess and OnFailure callbacks to get the operation result for the item.
@@ -56,7 +55,6 @@ type BulkIndexer interface {
 }
 
 // BulkIndexerConfig represents configuration of the indexer.
-//
 type BulkIndexerConfig struct {
 	NumWorkers    int           // The number of workers. Defaults to runtime.NumCPU().
 	FlushBytes    int           // The flush threshold in bytes. Defaults to 5MB.
@@ -88,7 +86,6 @@ type BulkIndexerConfig struct {
 }
 
 // BulkIndexerStats represents the indexer statistics.
-//
 type BulkIndexerStats struct {
 	NumAdded    uint64
 	NumFlushed  uint64
@@ -101,7 +98,6 @@ type BulkIndexerStats struct {
 }
 
 // BulkIndexerItem represents an indexer item.
-//
 type BulkIndexerItem struct {
 	Index           string
 	Action          string
@@ -111,21 +107,106 @@ type BulkIndexerItem struct {
 	VersionType     string
 	Body            io.Reader
 	RetryOnConflict *int
+	meta            bytes.Buffer // Item metadata header
+	payloadLength   int          // Item payload total length metadata+newline+body length
 
 	OnSuccess func(context.Context, BulkIndexerItem, BulkIndexerResponseItem)        // Per item
 	OnFailure func(context.Context, BulkIndexerItem, BulkIndexerResponseItem, error) // Per item
 }
 
+// marshallMeta format as JSON the item metadata.
+func (item *BulkIndexerItem) marshallMeta() {
+	var aux []byte
+	item.meta.WriteRune('{')
+	aux = strconv.AppendQuote(aux, item.Action)
+	item.meta.Write(aux)
+	aux = nil
+	item.meta.WriteRune(':')
+	item.meta.WriteRune('{')
+	if item.DocumentID != "" {
+		item.meta.WriteString(`"_id":`)
+		aux = strconv.AppendQuote(aux, item.DocumentID)
+		item.meta.Write(aux)
+		aux = nil
+	}
+
+	if item.DocumentID != "" && item.Version != nil {
+		item.meta.WriteRune(',')
+		item.meta.WriteString(`"version":`)
+		item.meta.WriteString(strconv.FormatInt(*item.Version, 10))
+	}
+
+	if item.DocumentID != "" && item.VersionType != "" {
+		item.meta.WriteRune(',')
+		item.meta.WriteString(`"version_type":`)
+		aux = strconv.AppendQuote(aux, item.VersionType)
+		item.meta.Write(aux)
+		aux = nil
+	}
+
+	if item.Routing != "" {
+		if item.DocumentID != "" {
+			item.meta.WriteRune(',')
+		}
+		item.meta.WriteString(`"routing":`)
+		aux = strconv.AppendQuote(aux, item.Routing)
+		item.meta.Write(aux)
+		aux = nil
+	}
+	if item.Index != "" {
+		if item.DocumentID != "" || item.Routing != "" {
+			item.meta.WriteRune(',')
+		}
+		item.meta.WriteString(`"_index":`)
+		aux = strconv.AppendQuote(aux, item.Index)
+		item.meta.Write(aux)
+		aux = nil
+	}
+	if item.RetryOnConflict != nil && item.Action == "update" {
+		if item.DocumentID != "" || item.Routing != "" || item.Index != "" {
+			item.meta.WriteString(",")
+		}
+		item.meta.WriteString(`"retry_on_conflict":`)
+		aux = strconv.AppendInt(aux, int64(*item.RetryOnConflict), 10)
+		item.meta.Write(aux)
+		aux = nil
+	}
+	item.meta.WriteRune('}')
+	item.meta.WriteRune('}')
+	item.meta.WriteRune('\n')
+}
+
+// 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
+		}
+
+		_, 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
+
+	return nil
+}
+
 // BulkIndexerResponse represents the Elasticsearch response.
-//
 type BulkIndexerResponse struct {
 	Took      int                                  `json:"took"`
 	HasErrors bool                                 `json:"errors"`
 	Items     []map[string]BulkIndexerResponseItem `json:"items,omitempty"`
 }
 
 // BulkIndexerResponseItem represents the Elasticsearch response item.
-//
 type BulkIndexerResponseItem struct {
 	Index      string `json:"_index"`
 	DocumentID string `json:"_id"`
@@ -152,13 +233,11 @@ type BulkIndexerResponseItem struct {
 }
 
 // BulkResponseJSONDecoder defines the interface for custom JSON decoders.
-//
 type BulkResponseJSONDecoder interface {
 	UnmarshalFromReader(io.Reader, *BulkIndexerResponse) error
 }
 
 // BulkIndexerDebugLogger defines the interface for a debugging logger.
-//
 type BulkIndexerDebugLogger interface {
 	Printf(string, ...interface{})
 }
@@ -186,7 +265,6 @@ type bulkIndexerStats struct {
 }
 
 // NewBulkIndexer creates a new bulk indexer.
-//
 func NewBulkIndexer(cfg BulkIndexerConfig) (BulkIndexer, error) {
 	if cfg.Client == nil {
 		cfg.Client, _ = elasticsearch.NewDefaultClient()
@@ -222,10 +300,16 @@ func NewBulkIndexer(cfg BulkIndexerConfig) (BulkIndexer, error) {
 // Add adds an item to the indexer.
 //
 // Adding an item after a call to Close() will panic.
-//
 func (bi *bulkIndexer) Add(ctx context.Context, item BulkIndexerItem) error {
 	atomic.AddUint64(&bi.stats.numAdded, 1)
 
+	// Serialize metadata to JSON
+	item.marshallMeta()
+	// Compute length for body & metadata
+	if err := item.computeLength(); err != nil {
+		return err
+	}
+
 	select {
 	case <-ctx.Done():
 		if bi.config.OnError != nil {
@@ -240,7 +324,6 @@ func (bi *bulkIndexer) Add(ctx context.Context, item BulkIndexerItem) error {
 
 // Close stops the periodic flush, closes the indexer queue channel,
 // notifies the done channel and calls flush on all writers.
-//
 func (bi *bulkIndexer) Close(ctx context.Context) error {
 	bi.ticker.Stop()
 	close(bi.queue)
@@ -273,7 +356,6 @@ func (bi *bulkIndexer) Close(ctx context.Context) error {
 }
 
 // Stats returns indexer statistics.
-//
 func (bi *bulkIndexer) Stats() BulkIndexerStats {
 	return BulkIndexerStats{
 		NumAdded:    atomic.LoadUint64(&bi.stats.numAdded),
@@ -288,7 +370,6 @@ func (bi *bulkIndexer) Stats() BulkIndexerStats {
 }
 
 // init initializes the bulk indexer.
-//
 func (bi *bulkIndexer) init() {
 	bi.queue = make(chan BulkIndexerItem, bi.config.NumWorkers)
 
@@ -334,7 +415,6 @@ func (bi *bulkIndexer) init() {
 }
 
 // worker represents an indexer worker.
-//
 type worker struct {
 	id    int
 	ch    <-chan BulkIndexerItem
@@ -346,7 +426,6 @@ type worker struct {
 }
 
 // run launches the worker in a goroutine.
-//
 func (w *worker) run() {
 	go func() {
 		ctx := context.Background()
@@ -363,7 +442,18 @@ func (w *worker) run() {
 				w.bi.config.DebugLogger.Printf("[worker-%03d] Received item [%s:%s]\n", w.id, item.Action, item.DocumentID)
 			}
 
-			if err := w.writeMeta(item); err != nil {
+			oversizePayload := w.bi.config.FlushBytes <= item.payloadLength
+			if !oversizePayload && w.buf.Len() > 0 && w.buf.Len()+item.payloadLength >= w.bi.config.FlushBytes {
+				if err := w.flush(ctx); err != nil {
+					w.mu.Unlock()
+					if w.bi.config.OnError != nil {
+						w.bi.config.OnError(ctx, err)
+					}
+					continue
+				}
+			}
+
+			if err := w.writeMeta(&item); err != nil {
 				if item.OnFailure != nil {
 					item.OnFailure(ctx, item, BulkIndexerResponseItem{}, err)
 				}
@@ -382,7 +472,11 @@ func (w *worker) run() {
 			}
 
 			w.items = append(w.items, item)
-			if w.buf.Len() >= w.bi.config.FlushBytes {
+			// Should the item payload exceed the configured FlushBytes flush happens instantly.
+			if oversizePayload {
+				if w.bi.config.DebugLogger != nil {
+					w.bi.config.DebugLogger.Printf("[worker-%03d] Oversize Payload in item [%s:%s]\n", w.id, item.Action, item.DocumentID)
+				}
 				if err := w.flush(ctx); err != nil {
 					w.mu.Unlock()
 					if w.bi.config.OnError != nil {
@@ -396,71 +490,15 @@ func (w *worker) run() {
 	}()
 }
 
-// writeMeta formats and writes the item metadata to the buffer; it must be called under a lock.
-//
-func (w *worker) writeMeta(item BulkIndexerItem) error {
-	w.buf.WriteRune('{')
-	w.aux = strconv.AppendQuote(w.aux, item.Action)
-	w.buf.Write(w.aux)
-	w.aux = w.aux[:0]
-	w.buf.WriteRune(':')
-	w.buf.WriteRune('{')
-	if item.DocumentID != "" {
-		w.buf.WriteString(`"_id":`)
-		w.aux = strconv.AppendQuote(w.aux, item.DocumentID)
-		w.buf.Write(w.aux)
-		w.aux = w.aux[:0]
+// writeMeta writes the item metadata to the buffer; it must be called under a lock.
+func (w *worker) writeMeta(item *BulkIndexerItem) error {
+	if _, err := w.buf.Write(item.meta.Bytes()); err != nil {
+		return err
 	}
-
-	if item.DocumentID != "" && item.Version != nil {
-		w.buf.WriteRune(',')
-		w.buf.WriteString(`"version":`)
-		w.buf.WriteString(strconv.FormatInt(*item.Version, 10))
-	}
-
-	if item.DocumentID != "" && item.VersionType != "" {
-		w.buf.WriteRune(',')
-		w.buf.WriteString(`"version_type":`)
-		w.aux = strconv.AppendQuote(w.aux, item.VersionType)
-		w.buf.Write(w.aux)
-		w.aux = w.aux[:0]
-	}
-
-	if item.Routing != "" {
-		if item.DocumentID != "" {
-			w.buf.WriteRune(',')
-		}
-		w.buf.WriteString(`"routing":`)
-		w.aux = strconv.AppendQuote(w.aux, item.Routing)
-		w.buf.Write(w.aux)
-		w.aux = w.aux[:0]
-	}
-	if item.Index != "" {
-		if item.DocumentID != "" || item.Routing != "" {
-			w.buf.WriteRune(',')
-		}
-		w.buf.WriteString(`"_index":`)
-		w.aux = strconv.AppendQuote(w.aux, item.Index)
-		w.buf.Write(w.aux)
-		w.aux = w.aux[:0]
-	}
-	if item.RetryOnConflict != nil && item.Action == "update" {
-		if item.DocumentID != "" || item.Routing != "" || item.Index != "" {
-			w.buf.WriteString(",")
-		}
-		w.buf.WriteString(`"retry_on_conflict":`)
-		w.aux = strconv.AppendInt(w.aux, int64(*item.RetryOnConflict), 10)
-		w.buf.Write(w.aux)
-		w.aux = w.aux[:0]
-	}
-	w.buf.WriteRune('}')
-	w.buf.WriteRune('}')
-	w.buf.WriteRune('\n')
 	return nil
 }
 
 // writeBody writes the item body to the buffer; it must be called under a lock.
-//
 func (w *worker) writeBody(item *BulkIndexerItem) error {
 	if item.Body != nil {
 
@@ -492,7 +530,6 @@ func (w *worker) writeBody(item *BulkIndexerItem) error {
 }
 
 // flush writes out the worker buffer; it must be called under a lock.
-//
 func (w *worker) flush(ctx context.Context) error {
 	if w.bi.config.OnFlushStart != nil {
 		ctx = w.bi.config.OnFlushStart(ctx)
@@ -515,8 +552,12 @@ func (w *worker) flush(ctx context.Context) error {
 	)
 
 	defer func() {
-		w.items = w.items[:0]
-		w.buf.Reset()
+		w.items = nil
+		if w.buf.Cap() > w.bi.config.FlushBytes {
+			w.buf = bytes.NewBuffer(make([]byte, 0, w.bi.config.FlushBytes))
+		} else {
+			w.buf.Reset()
+		}
 	}()
 
 	if w.bi.config.DebugLogger != nil {

esutil/bulk_indexer_integration_test.go

@@ -72,7 +72,6 @@ func TestBulkIndexerIntegration(t *testing.T) {
 				bi, _ := esutil.NewBulkIndexer(esutil.BulkIndexerConfig{
 					Index:  indexName,
 					Client: es,
-					// FlushBytes: 3e+6,
 				})
 
 				numItems := 100000