fix: BulkIndexer NumAdded overcount on cancelled Add() (#784)

fix(opensearchutil): correct BulkIndexerStats.NumAdded overcount on cancelled context (#784)

NumAdded was incremented unconditionally in Add(), even when the caller's
context was already cancelled and the item never reached the queue. Move
the increment after the successful queue send and introduce
BulkAddFailCount on the <-ctx.Done() branch so callers can distinguish
queued items from items the indexer rejected.

Migrate every bulkIndexerStats counter to sync/atomic.Uint64 typed values
so future direct access fails to compile rather than only surfacing under
`-race`.

Add a unit test that verifies, with an already-cancelled context, every
Add() call ends up in exactly one of NumAdded or BulkAddFailCount.

Fixes #783

Co-authored-by: evgenigourvitch <evgeni.evgeni@gmail.com>
Signed-off-by: evgenigourvitch <evgeni.evgeni@gmail.com>
Signed-off-by: Sean Chittenden <sean.chittenden@crowdstrike.com>

Author: evgenigourvitch

CHANGELOG.md

@@ -208,6 +208,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 
 ### Fixed
 
+- Fix `BulkIndexerStats.NumAdded` overcounting items rejected by `Add()` when the caller's context is cancelled before the item could be enqueued: increment `NumAdded` only after the queue accepts the item, and add a new `BulkAddFailCount` counter for items dropped on the `<-ctx.Done()` branch. Migrate `bulkIndexerStats` fields to `sync/atomic.Uint64` typed values so future direct access is a compile-time error rather than a `-race`-only finding ([#783](https://github.com/opensearch-project/opensearch-go/issues/783))
 - Add typed response-format defaults for `v5preview/opensearchapi/` cat, list, ppl, and sql operations: when the caller leaves `Format` unset, the SDK now emits the value the typed Resp struct expects (`json` for cat/list/explain, `jdbc` for ppl/sql query) instead of letting the server fall back to a default the JSON decoder cannot handle.
 - Replace `WaitForAllNodesReady` inline `require.Eventually` loop with a layered readiness FSM (`internal/test/readiness`) that observes per-node progression through `LayerTCP -> LayerHTTP -> LayerClusterJoin -> LayerStatsReady`, records transitions including regressions, and emits a structured per-node diagnostic with the full last cat-nodes response on timeout. Per-layer budgets are tuned for CI pessimism (cold JVM startup is the long pole); total budget for `TargetClusterReady` is 6.5 minutes. ([#650](https://github.com/opensearch-project/opensearch-go/issues/650))
 - Fix bulk indexer HTML-escaping `_id` and `routing` values containing `<`, `>`, or `&` characters, causing OpenSearch to store escaped values (e.g., `\u003croot_account\u003e` stored instead of `<root_account>`), leading to duplicate documents, unreachable data on read-by-ID paths, and potential shard routing mismatches. Present since the `json.Marshal` migration in 2021 (commit `3da59092`). Replace `json.Marshal` with `json.NewEncoder` + `SetEscapeHTML(false)` in `opensearchutil.worker.writeMeta` and `opensearchutil.JSONReader`; replace per-worker `aux []byte` with `sync.Pool`-backed `*bytes.Buffer`; add table-driven test coverage for `writeMeta` edge cases and refactor remaining `TestBulkIndexer` subtests to table-driven `require`-based style ([#824](https://github.com/opensearch-project/opensearch-go/pull/824))

opensearchutil/bulk_indexer.go

@@ -105,14 +105,15 @@ type BulkIndexerConfig struct {
 
 // BulkIndexerStats represents the indexer statistics.
 type BulkIndexerStats struct {
-	NumAdded    uint64
-	NumFlushed  uint64
-	NumFailed   uint64
-	NumIndexed  uint64
-	NumCreated  uint64
-	NumUpdated  uint64
-	NumDeleted  uint64
-	NumRequests uint64
+	NumAdded         uint64
+	BulkAddFailCount uint64 // Items rejected by Add() because the caller's context was cancelled before the item could be enqueued.
+	NumFlushed       uint64
+	NumFailed        uint64
+	NumIndexed       uint64
+	NumCreated       uint64
+	NumUpdated       uint64
+	NumDeleted       uint64
+	NumRequests      uint64
 }
 
 // BulkIndexerItem represents an indexer item.
@@ -169,14 +170,15 @@ type bulkIndexer struct {
 }
 
 type bulkIndexerStats struct {
-	numAdded    uint64
-	numFlushed  uint64
-	numFailed   uint64
-	numIndexed  uint64
-	numCreated  uint64
-	numUpdated  uint64
-	numDeleted  uint64
-	numRequests uint64
+	numAdded         atomic.Uint64
+	bulkAddFailCount atomic.Uint64
+	numFlushed       atomic.Uint64
+	numFailed        atomic.Uint64
+	numIndexed       atomic.Uint64
+	numCreated       atomic.Uint64
+	numUpdated       atomic.Uint64
+	numDeleted       atomic.Uint64
+	numRequests      atomic.Uint64
 }
 
 // NewBulkIndexer creates a new bulk indexer.
@@ -232,15 +234,15 @@ func NewBulkIndexer(cfg BulkIndexerConfig) (BulkIndexer, error) {
 //
 // 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)
-
 	select {
 	case <-ctx.Done():
+		bi.stats.bulkAddFailCount.Add(1)
 		if bi.config.OnError != nil {
 			bi.config.OnError(ctx, ctx.Err())
 		}
 		return ctx.Err()
 	case bi.queue <- item:
+		bi.stats.numAdded.Add(1)
 	}
 
 	return nil
@@ -283,14 +285,15 @@ 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),
-		NumFlushed:  atomic.LoadUint64(&bi.stats.numFlushed),
-		NumFailed:   atomic.LoadUint64(&bi.stats.numFailed),
-		NumIndexed:  atomic.LoadUint64(&bi.stats.numIndexed),
-		NumCreated:  atomic.LoadUint64(&bi.stats.numCreated),
-		NumUpdated:  atomic.LoadUint64(&bi.stats.numUpdated),
-		NumDeleted:  atomic.LoadUint64(&bi.stats.numDeleted),
-		NumRequests: atomic.LoadUint64(&bi.stats.numRequests),
+		NumAdded:         bi.stats.numAdded.Load(),
+		BulkAddFailCount: bi.stats.bulkAddFailCount.Load(),
+		NumFlushed:       bi.stats.numFlushed.Load(),
+		NumFailed:        bi.stats.numFailed.Load(),
+		NumIndexed:       bi.stats.numIndexed.Load(),
+		NumCreated:       bi.stats.numCreated.Load(),
+		NumUpdated:       bi.stats.numUpdated.Load(),
+		NumDeleted:       bi.stats.numDeleted.Load(),
+		NumRequests:      bi.stats.numRequests.Load(),
 	}
 }
 
@@ -388,7 +391,7 @@ func (w *worker) run(ctx context.Context) {
 						item.OnFailure(ctx, item, opensearchapi.BulkRespItem{}, err)
 					}
 
-					atomic.AddUint64(&w.bi.stats.numFailed, 1)
+					w.bi.stats.numFailed.Add(1)
 					w.mu.Unlock()
 
 					continue
@@ -398,7 +401,7 @@ func (w *worker) run(ctx context.Context) {
 					if item.OnFailure != nil {
 						item.OnFailure(ctx, item, opensearchapi.BulkRespItem{}, err)
 					}
-					atomic.AddUint64(&w.bi.stats.numFailed, 1)
+					w.bi.stats.numFailed.Add(1)
 					w.mu.Unlock()
 
 					continue
@@ -528,7 +531,7 @@ func (w *worker) flush(ctx context.Context) error {
 		w.bi.config.DebugLogger.Printf("[worker-%03d] Flush: %s\n", w.id, w.buf.String())
 	}
 
-	atomic.AddUint64(&w.bi.stats.numRequests, 1)
+	w.bi.stats.numRequests.Add(1)
 	req := opensearchapi.BulkReq{
 		Index: w.bi.config.Index,
 		Body:  w.buf,
@@ -576,22 +579,22 @@ func (w *worker) flush(ctx context.Context) error {
 			info = v
 		}
 		if info.Error != nil || info.Status >= http.StatusMultipleChoices {
-			atomic.AddUint64(&w.bi.stats.numFailed, 1)
+			w.bi.stats.numFailed.Add(1)
 			if item.OnFailure != nil {
 				item.OnFailure(ctx, item, info, nil)
 			}
 		} else {
-			atomic.AddUint64(&w.bi.stats.numFlushed, 1)
+			w.bi.stats.numFlushed.Add(1)
 
 			switch op {
 			case "index":
-				atomic.AddUint64(&w.bi.stats.numIndexed, 1)
+				w.bi.stats.numIndexed.Add(1)
 			case "create":
-				atomic.AddUint64(&w.bi.stats.numCreated, 1)
+				w.bi.stats.numCreated.Add(1)
 			case "delete":
-				atomic.AddUint64(&w.bi.stats.numDeleted, 1)
+				w.bi.stats.numDeleted.Add(1)
 			case "update":
-				atomic.AddUint64(&w.bi.stats.numUpdated, 1)
+				w.bi.stats.numUpdated.Add(1)
 			}
 
 			if item.OnSuccess != nil {
@@ -604,7 +607,7 @@ func (w *worker) flush(ctx context.Context) error {
 }
 
 func (w *worker) handleBulkError(ctx context.Context, err error) error {
-	atomic.AddUint64(&w.bi.stats.numFailed, uint64(len(w.items)))
+	w.bi.stats.numFailed.Add(uint64(len(w.items)))
 
 	// info (the response item) will be empty since the bulk request failed
 	var info opensearchapi.BulkRespItem

opensearchutil/bulk_indexer_internal_test.go

@@ -499,6 +499,38 @@ func TestBulkIndexerContext(t *testing.T) {
 				require.True(t, gotDeadline, "expected at least one context.DeadlineExceeded in: %q", errs)
 			},
 		},
+		{
+			name: "Add does not increment NumAdded when context is already cancelled",
+			run: func(t *testing.T) {
+				client, _ := opensearchapi.NewClient(opensearchapi.Config{Client: opensearch.Config{Transport: &mockTransport{}}})
+				bi, _ := NewBulkIndexer(BulkIndexerConfig{NumWorkers: 1, Client: client})
+
+				ctx, cancel := context.WithCancel(context.Background())
+				cancel()
+
+				// select{} chooses randomly between the two ready cases when the
+				// queue still has room, so we cannot assert every Add fails. We can
+				// assert the bookkeeping invariant: each Add ends up in exactly one
+				// of NumAdded or BulkAddFailCount, never both, never neither.
+				const numAttempts = 50
+				var nilReturns, errReturns uint64
+				for range numAttempts {
+					if err := bi.Add(ctx, BulkIndexerItem{Action: "index", DocumentID: "cancelled"}); err == nil {
+						nilReturns++
+					} else {
+						require.ErrorIs(t, err, context.Canceled)
+						errReturns++
+					}
+				}
+				require.NoError(t, bi.Close(context.Background()))
+
+				stats := bi.Stats()
+				require.Equal(t, nilReturns, stats.NumAdded, "NumAdded must equal the number of Add() calls that returned nil")
+				require.Equal(t, errReturns, stats.BulkAddFailCount, "BulkAddFailCount must equal the number of Add() calls that returned ctx.Err()")
+				require.Equal(t, uint64(numAttempts), stats.NumAdded+stats.BulkAddFailCount, "every Add() must be accounted for exactly once")
+				require.Greater(t, errReturns, uint64(0), "at least one Add() should fail when context is already cancelled")
+			},
+		},
 		{
 			name: "Close returns error on cancelled context",
 			run: func(t *testing.T) {