Fix singleton stack-corruption NPE in DatetimeUdtNormalizeRule (opensearch-project#5458)

* fix: instantiate DatetimeUdt normalize/output-cast rules per plan() call

DatetimeUdtNormalizeRule and DatetimeOutputCastRule extend
RelHomogeneousShuttle, which inherits a stateful Deque<RelNode> stack
from RelShuttleImpl. DatetimeExtension.postAnalysisRules() returned
the static INSTANCE of each rule, sharing the same shuttle (and the
same stack) across every UnifiedQueryPlanner.plan() invocation.

If any traversal ever ends with an unbalanced stack, residual entries
persist to the next query. The next query's visitChild() then pops a
stale or empty stack and throws NoSuchElementException at
RelShuttleImpl.visitChild line 67 (the stack.pop() in the finally
block) — surfacing as the cluster-side stack trace reported on
analytics-engine-routed parquet indices for queries that combine
aggregations over datetime UDT columns (e.g.
"stats count() as field_count, distinct_count(field)").

Return fresh instances per plan() instead. Drop the INSTANCE
constants and the Lombok @NoArgsConstructor on both rules; document
the singleton-unsafety on each class JavaDoc.

Add a regression test that runs several plan() calls in sequence
against the same context, covering stats+distinct_count over both
schema-declared and eval-derived datetime columns.

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

* test: add analytics-engine regression IT for singleton stack-corruption

CalciteDatetimeUdtNormalizeRegressionIT exercises the failure pattern
that triggered the cluster-side NoSuchElementException:
stats + distinct_count over datetime columns, repeated 20 times to
amplify any plan() carry-over.

The IT is harness-aware:
- Without `-Dtests.analytics.force_routing=true`: queries go through
  the V2 / Calcite engine path. The DatetimeUdtNormalizeRule path is
  not exercised, so the IT passes as a baseline correctness check.
- With `-Dtests.analytics.force_routing=true
  -Dtests.analytics.parquet_indices=true`: every query routes through
  the analytics-engine path and hits the DatetimeUdtNormalizeRule
  shuttle that this PR fixes. The 20-iteration pattern surfaces any
  remaining singleton-stack carry-over.

CI's :integTest task (in-process testCluster without analytics-engine)
runs the IT through the V2 path, which is safe and fast. The
analytics-engine verification path is via :integTestRemote against an
externally-managed cluster built per
`docs/dev/ppl-analytics-engine-routing.md`.

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

* 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>

* test: rename to CalcitePlannerConcurrencyIT (review nit)

@dai-chen flagged that the IT name was over-scoped to a single rule and
the file would read better as a general bucket for planner-level
concurrency / state-isolation regressions. The actual surface under test
is UnifiedQueryPlanner's post-analysis pipeline — any RelShuttle
extension that doesn't isolate per-call state is unsafe under concurrent
load, not just the datetime rules.

Renames the file and class, updates the JavaDoc to describe the planner-
level invariant rather than the specific Datetime* rules, and notes the
current cases as the regression that motivated the suite. Test method
bodies and assertions are unchanged.

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

---------

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

Author: ahkcs

api/src/main/java/org/opensearch/sql/api/spec/datetime/DatetimeExtension.java

@@ -22,7 +22,8 @@ public class DatetimeExtension implements LanguageExtension {
 
   @Override
   public List<RelShuttle> postAnalysisRules() {
-    return List.of(DatetimeUdtNormalizeRule.INSTANCE, DatetimeOutputCastRule.INSTANCE);
+    // Fresh instances per plan() because RelHomogeneousShuttle inherits a stateful stack.
+    return List.of(new DatetimeUdtNormalizeRule(), new DatetimeOutputCastRule());
   }
 
   /** Maps datetime UDT types to their standard Calcite equivalents. */

api/src/main/java/org/opensearch/sql/api/spec/datetime/DatetimeOutputCastRule.java

@@ -9,8 +9,6 @@
 
 import java.util.ArrayList;
 import java.util.List;
-import lombok.AccessLevel;
-import lombok.NoArgsConstructor;
 import org.apache.calcite.rel.RelHomogeneousShuttle;
 import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.logical.LogicalProject;
@@ -21,12 +19,14 @@
 import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.sql.type.SqlTypeName;
 
-/** Wraps the root output with CAST(datetime → VARCHAR) for PPL wire-format compatibility. */
-@NoArgsConstructor(access = AccessLevel.PRIVATE)
+/**
+ * Wraps the root output with CAST(datetime → VARCHAR) for PPL wire-format compatibility.
+ *
+ * <p>Not a singleton: {@link RelHomogeneousShuttle} inherits a stateful {@code stack} field from
+ * {@link org.apache.calcite.rel.RelShuttleImpl}, so a fresh instance must be used per plan().
+ */
 class DatetimeOutputCastRule extends RelHomogeneousShuttle {
 
-  static final DatetimeOutputCastRule INSTANCE = new DatetimeOutputCastRule();
-
   @Override
   public RelNode visit(RelNode other) {
     List<RelDataTypeField> fields = other.getRowType().getFieldList();

api/src/main/java/org/opensearch/sql/api/spec/datetime/DatetimeUdtNormalizeRule.java

@@ -6,8 +6,6 @@
 package org.opensearch.sql.api.spec.datetime;
 
 import java.util.Optional;
-import lombok.AccessLevel;
-import lombok.NoArgsConstructor;
 import org.apache.calcite.rel.RelHomogeneousShuttle;
 import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.type.RelDataType;
@@ -22,12 +20,12 @@
 /**
  * Temporary patch that rewrites datetime UDT return types on RexCall nodes to standard Calcite
  * types.
+ *
+ * <p>Not a singleton: {@link RelHomogeneousShuttle} inherits a stateful {@code stack} field from
+ * {@link org.apache.calcite.rel.RelShuttleImpl}, so a fresh instance must be used per plan().
  */
-@NoArgsConstructor(access = AccessLevel.PRIVATE)
 class DatetimeUdtNormalizeRule extends RelHomogeneousShuttle {
 
-  static final DatetimeUdtNormalizeRule INSTANCE = new DatetimeUdtNormalizeRule();
-
   @Override
   public RelNode visit(RelNode other) {
     RelNode visited = super.visit(other);

api/src/test/java/org/opensearch/sql/api/spec/datetime/DatetimeExtensionTest.java

@@ -154,6 +154,27 @@ public void testNonDatetimeFieldsNotWrapped() {
             """);
   }
 
+  @Test
+  public void testSequentialPlanCallsDoNotCorruptShuttleStack() {
+    // Regression test: DatetimeUdtNormalizeRule extends RelHomogeneousShuttle which inherits a
+    // stateful Deque<RelNode> stack field. Earlier implementations used a static INSTANCE shared
+    // across all plan() calls; under workloads with aggregations (especially count + distinct_count
+    // over datetime columns), the shared stack would desynchronize and visitChild's pop would throw
+    // NoSuchElementException on subsequent plan calls. Running several distinct plans through the
+    // same context confirms each invocation gets a fresh shuttle.
+    for (int i = 0; i < 5; i++) {
+      planner.plan(
+          "source = catalog.events"
+              + " | stats count() as field_count, distinct_count(created_at) as distinct_count");
+      planner.plan(
+          "source = catalog.events"
+              + " | eval ts = TIMESTAMP(name)"
+              + " | stats count() as field_count, distinct_count(ts) as distinct_count");
+      planner.plan(
+          "source = catalog.events | where created_at > \"2024-01-01\" | fields hire_date");
+    }
+  }
+
   @Test
   public void testOutputCastCanCompileAndExecute() throws Exception {
     RelNode plan =

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

@@ -0,0 +1,155 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.calcite.remote;
+
+import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_DATE_FORMATS;
+
+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;
+
+/**
+ * Integration tests for {@code UnifiedQueryPlanner} state isolation under concurrent load.
+ *
+ * <p>The planner's post-analysis pipeline (extensions registered via {@code
+ * LanguageSpec.postAnalysisRules}) uses Calcite {@code RelShuttle} subclasses. {@code
+ * RelShuttleImpl} inherits a non-thread-safe {@code ArrayDeque<RelNode>} stack used by {@code
+ * visitChild}'s push/pop. Any extension that returns the same shuttle instance across {@code
+ * plan()} calls is unsafe under concurrent load: 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 test methods here fire many queries through a thread pool to exercise concurrent {@code
+ * plan()} invocations. New planner-level concurrency / state-isolation regressions belong in this
+ * class. The current cases cover {@code DatetimeExtension}'s {@code RelHomogeneousShuttle}
+ * subclasses ({@code DatetimeUdtNormalizeRule}, {@code DatetimeOutputCastRule}) which were
+ * previously returned as static {@code INSTANCE}s and caused the production failure that motivated
+ * this suite.
+ *
+ * <p>Run via:
+ *
+ * <pre>{@code
+ * ./gradlew :integ-test:integTestRemote \
+ *   -Dtests.rest.cluster=localhost:9200 -Dtests.cluster=localhost:9300 \
+ *   -Dtests.clustername=runTask \
+ *   -Dtests.analytics.force_routing=true \
+ *   -Dtests.analytics.parquet_indices=true \
+ *   --tests org.opensearch.sql.calcite.remote.CalcitePlannerConcurrencyIT
+ * }</pre>
+ */
+public class CalcitePlannerConcurrencyIT extends PPLIntegTestCase {
+
+  /** 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 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 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 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 planner's post-analysis shuttles
+    // through different visitChild call counts — amplifying any cross-query stack pollution.
+    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);
+  }
+
+  /**
+   * Fire all queries through a fixed-size thread pool. Asserts every query completes without
+   * exception. With a shuttle-state-leak 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);
+    }
+  }
+}