test: switch regression IT to concurrent query pattern

The sequential iteration variant passed even with the singleton bug in
place — local cluster doesn't carry over enough state between calls in
one thread. The actual production trigger is parallel queries from a
dashboard "field statistics" panel: multiple cluster threads call
plan() simultaneously, all using the shared singleton's non-thread-safe
ArrayDeque. Their push/pop operations interleave and corrupt the stack.

Verified locally against analytics-engine path with parquet indices:
- Unfixed cluster: 2-3 / 80 queries fail with NoSuchElementException
  (HTTP 500), matching the production stack trace exactly.
- Fixed cluster: 0 / 80 failures.

Uses CompletableFuture + 8-thread pool to fire 80 queries per test
across:
- testConcurrentStatsDistinctCountOverDatetime: same shape, varied
  datetime fields.
- testConcurrentMixedDatetimePlans: three different plan shapes
  interleaved — mixed visitChild call counts amplify the race.

Signed-off-by: Kai Huang <ahkcs@amazon.com>

Author: ahkcs

integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteDatetimeUdtNormalizeRegressionIT.java

@@ -5,30 +5,34 @@
 
 package org.opensearch.sql.calcite.remote;
 
-import static org.junit.Assert.assertNotNull;
 import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_DATE_FORMATS;
 
-import java.io.IOException;
-import org.json.JSONObject;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
 import org.junit.jupiter.api.Test;
 import org.opensearch.sql.ppl.PPLIntegTestCase;
 
 /**
  * Regression IT for the {@code DatetimeUdtNormalizeRule} / {@code DatetimeOutputCastRule}
  * singleton-stack-corruption bug.
  *
- * <p>Both rules extend Calcite's {@code RelHomogeneousShuttle}, which inherits a stateful {@code
- * Deque<RelNode>} stack from {@code RelShuttleImpl}. Earlier code returned the same {@code
- * INSTANCE} of each rule from {@code DatetimeExtension.postAnalysisRules()} on every {@code
- * UnifiedQueryPlanner.plan()} call. If any traversal ever finished with an unbalanced stack,
- * residual entries persisted to the next query and the next {@code visitChild()} popped a stale or
- * empty stack — surfacing as {@code NoSuchElementException} at {@code RelShuttleImpl.visitChild}
- * line 67.
+ * <p>Both rules extend Calcite's {@code RelHomogeneousShuttle}, which inherits a stateful
+ * non-thread-safe {@code ArrayDeque<RelNode>} stack from {@code RelShuttleImpl}. Earlier code
+ * returned the same {@code INSTANCE} of each rule from {@code
+ * DatetimeExtension.postAnalysisRules()} on every {@code UnifiedQueryPlanner.plan()} call. Under
+ * parallel query load — exactly what a dashboard "field statistics" panel issues when it probes
+ * every field in an index concurrently — multiple cluster threads call {@code plan()}
+ * simultaneously and their push/pop on the shared stack interleave, leaving residual entries that
+ * surface on a subsequent traversal as {@code NoSuchElementException} at {@code
+ * RelShuttleImpl.visitChild} line 67 (the {@code stack.pop()} in the {@code finally} block).
  *
- * <p>The failure was reported as intermittent on dashboards issuing field-statistics queries of the
- * shape {@code stats count() as field_count, distinct_count(field)} against parquet-backed indices.
- * This IT runs the same query shape repeatedly against a parquet-backed (composite) index with
- * multiple datetime columns to exercise the analytics-engine route end-to-end.
+ * <p>This IT reproduces the production failure by firing many {@code distinct_count} queries over
+ * datetime fields concurrently against a parquet-backed (composite) index.
  *
  * <p>Run via:
  *
@@ -43,75 +47,107 @@
  */
 public class CalciteDatetimeUdtNormalizeRegressionIT extends PPLIntegTestCase {
 
-  /** Number of repetitions per test. Singleton bug surfaces order-dependent across plan() calls. */
-  private static final int ITERATIONS = 20;
+  /** Concurrency level — matches the rough parallelism of a dashboard field-stats panel. */
+  private static final int PARALLELISM = 8;
+
+  /** Total queries fired per test. */
+  private static final int QUERIES = 80;
 
   @Override
   public void init() throws Exception {
     super.init();
     enableCalcite();
 
-    // DATE_FORMATS index has many date columns of different formats / precisions, which is what
-    // dashboard field-statistics panels iterate over. Loaded through the helper so that with
-    // -Dtests.analytics.parquet_indices=true it gets provisioned as a parquet-backed composite
-    // index — required for analytics-engine routing.
+    // DATE_FORMATS has many datetime columns of different formats/precisions. With
+    // -Dtests.analytics.parquet_indices=true the helper provisions it as a parquet-backed
+    // composite index — required for analytics-engine routing.
     loadIndex(Index.DATE_FORMATS);
   }
 
   @Test
-  public void testSequentialStatsDistinctCountOverDatetime() throws IOException {
-    // Bug repro: each iteration runs a stats + distinct_count over a datetime field. With the
-    // singleton rule, residual entries on the shuttle's internal stack from the previous plan()
-    // would cause NoSuchElementException on a subsequent traversal. With fresh instances per
-    // plan(), the stack is always empty at entry and the iterations all succeed.
-    String query =
-        String.format(
-            "source=%s | stats count() as field_count, distinct_count(epoch_millis) as"
-                + " distinct_count",
-            TEST_INDEX_DATE_FORMATS);
-    for (int i = 0; i < ITERATIONS; i++) {
-      JSONObject result = executeQuery(query);
-      // Sanity: response was returned without a cluster-side exception.
-      assertNotNull("iteration " + i + " produced no result", result);
+  public void testConcurrentStatsDistinctCountOverDatetime() throws Exception {
+    String[] fields = {
+      "epoch_millis", "epoch_second", "date_optional_time", "strict_date_optional_time"
+    };
+    List<String> queries = new ArrayList<>(QUERIES);
+    for (int i = 0; i < QUERIES; i++) {
+      String field = fields[i % fields.length];
+      queries.add(
+          String.format(
+              "source=%s | stats count() as field_count, distinct_count(%s) as distinct_count",
+              TEST_INDEX_DATE_FORMATS, field));
     }
+    executeConcurrent(queries);
   }
 
   @Test
-  public void testInterleavedStatsAndDatetimeProjection() throws IOException {
-    // Bug repro variant: interleave stats+distinct_count with a plain datetime projection that
-    // exercises the DatetimeOutputCastRule. Different plan shapes per iteration push the rule
-    // through different visitChild paths and surface stack desync faster.
-    for (int i = 0; i < ITERATIONS; i++) {
-      executeQuery(
-          String.format(
-              "source=%s | stats count() as field_count, distinct_count(epoch_millis) as"
-                  + " distinct_count",
-              TEST_INDEX_DATE_FORMATS));
-      executeQuery(
-          String.format(
-              "source=%s | fields epoch_millis, epoch_second, date_optional_time",
-              TEST_INDEX_DATE_FORMATS));
-      executeQuery(
-          String.format(
-              "source=%s | stats count() as field_count, distinct_count(epoch_second) as"
-                  + " distinct_count",
-              TEST_INDEX_DATE_FORMATS));
+  public void testConcurrentMixedDatetimePlans() throws Exception {
+    // Mix three plan shapes: stats+distinct_count, plain field projection (datetime cast), and
+    // stats by a different field. Different plan shapes push the singleton shuttle through
+    // different visitChild call counts — making cross-query stack pollution more likely.
+    List<String> shapes =
+        List.of(
+            "source=%s | stats count() as field_count, distinct_count(epoch_millis) as"
+                + " distinct_count",
+            "source=%s | fields epoch_millis, epoch_second, date_optional_time",
+            "source=%s | stats count() as field_count, distinct_count(epoch_second) as"
+                + " distinct_count by date_optional_time");
+    List<String> queries = new ArrayList<>(QUERIES);
+    for (int i = 0; i < QUERIES; i++) {
+      queries.add(String.format(shapes.get(i % shapes.size()), TEST_INDEX_DATE_FORMATS));
     }
+    executeConcurrent(queries);
   }
 
-  @Test
-  public void testDistinctCountOverMultipleDatetimeFields() throws IOException {
-    // Bug repro variant: iterate distinct_count over different datetime fields in sequence —
-    // mirrors the dashboard field-statistics tab that probes every field in the index.
-    String[] datetimeFields = {
-      "epoch_millis", "epoch_second", "date_optional_time", "strict_date_optional_time"
-    };
-    for (int i = 0; i < ITERATIONS; i++) {
-      String field = datetimeFields[i % datetimeFields.length];
-      executeQuery(
-          String.format(
-              "source=%s | stats count() as field_count, distinct_count(%s) as distinct_count",
-              TEST_INDEX_DATE_FORMATS, field));
+  /**
+   * Fire all queries through a fixed-size thread pool. Asserts every query completes without
+   * exception. With the singleton bug present this triggers {@code NoSuchElementException} on at
+   * least one task once the stack interleaving corrupts state.
+   */
+  private void executeConcurrent(List<String> queries) throws Exception {
+    var executor = Executors.newFixedThreadPool(PARALLELISM);
+    try {
+      List<CompletableFuture<Void>> futures = new ArrayList<>(queries.size());
+      AtomicInteger failures = new AtomicInteger();
+      List<Throwable> errors = new ArrayList<>();
+      for (String query : queries) {
+        futures.add(
+            CompletableFuture.runAsync(
+                () -> {
+                  try {
+                    executeQuery(query);
+                  } catch (Exception e) {
+                    failures.incrementAndGet();
+                    synchronized (errors) {
+                      errors.add(e);
+                    }
+                  }
+                },
+                executor));
+      }
+      for (CompletableFuture<Void> f : futures) {
+        try {
+          f.get(60, TimeUnit.SECONDS);
+        } catch (ExecutionException e) {
+          failures.incrementAndGet();
+          synchronized (errors) {
+            errors.add(e.getCause());
+          }
+        }
+      }
+      if (failures.get() > 0) {
+        StringBuilder msg = new StringBuilder();
+        msg.append(failures.get()).append("/").append(queries.size()).append(" queries failed:");
+        synchronized (errors) {
+          for (int i = 0; i < Math.min(3, errors.size()); i++) {
+            msg.append("\n  - ").append(errors.get(i));
+          }
+        }
+        throw new AssertionError(msg.toString());
+      }
+    } finally {
+      executor.shutdown();
+      executor.awaitTermination(30, TimeUnit.SECONDS);
     }
   }
 }