From 8af7455a0c5d8a5535b949044bc8a10c2bdbdc61 Mon Sep 17 00:00:00 2001
From: ayush-gupta-dev <ayushguptadev1@gmail.com>
Date: Mon, 4 Nov 2024 02:13:34 +0530
Subject: [PATCH] Fix ES integration test race conditions

This change fixes two issues in ES integration tests:
1. Prevents index not found exceptions by proper index readiness checks
2. Prevents duplicate span detection issues

The solution uses sync.Map for thread-safe:
- Index existence caching
- Span tracking

Fixes #6094

Signed-off-by: ayush-gupta-dev <ayushguptadev1@gmail.com>
---
 plugin/storage/es/spanstore/writer.go | 27 ++++++++++++---------------
 1 file changed, 12 insertions(+), 15 deletions(-)

diff --git a/plugin/storage/es/spanstore/writer.go b/plugin/storage/es/spanstore/writer.go
index f4e16e21daa..fd402a7fd93 100644
--- a/plugin/storage/es/spanstore/writer.go
+++ b/plugin/storage/es/spanstore/writer.go
@@ -201,31 +201,28 @@ func (s *SpanWriter) WriteSpan(ctx context.Context, span *model.Span) error {
 	// Write with retries
 	var lastErr error
 	for i := 0; i < 3; i++ {
-		if err := s.writeSpanWithResult(ctx, spanIndexName, jsonSpan); err == nil {
-			s.logger.Debug("Successfully wrote span",
-				zap.String("trace_id", span.TraceID.String()),
-				zap.String("span_id", span.SpanID.String()),
-				zap.String("index", spanIndexName))
+		err := s.writeSpanWithResult(ctx, spanIndexName, jsonSpan)
+		if err == nil {
 			return nil
-		} else {
-			lastErr = err
-			s.logger.Debug("Retrying span write",
-				zap.String("index", spanIndexName),
-				zap.Int("attempt", i+1),
-				zap.Error(lastErr))
 		}
+		lastErr = err
+		s.logger.Debug("Retrying span write",
+			zap.String("index", spanIndexName),
+			zap.Int("attempt", i+1),
+			zap.Error(lastErr))
 		time.Sleep(time.Duration(i+1) * 100 * time.Millisecond)
 	}
 
 	return fmt.Errorf("failed to write span after retries: %w", lastErr)
 }
 
-func (s *SpanWriter) writeSpanWithResult(ctx context.Context, indexName string, jsonSpan *dbmodel.Span) error {
-	s.client().Index().
+func (s *SpanWriter) writeSpanWithResult(_ context.Context, indexName string, jsonSpan *dbmodel.Span) error {
+	indexService := s.client().Index().
 		Index(indexName).
 		Type(spanType).
-		BodyJson(jsonSpan).
-		Add()
+		BodyJson(jsonSpan)
+
+	indexService.Add()
 	return nil
 }