diff --git a/api/src/main/java/org/opensearch/sql/api/spec/datetime/DatetimeExtension.java b/api/src/main/java/org/opensearch/sql/api/spec/datetime/DatetimeExtension.java index 944ac4a4bf1..1f0b0b820a5 100644 --- a/api/src/main/java/org/opensearch/sql/api/spec/datetime/DatetimeExtension.java +++ b/api/src/main/java/org/opensearch/sql/api/spec/datetime/DatetimeExtension.java @@ -22,7 +22,8 @@ public class DatetimeExtension implements LanguageExtension { @Override public List 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. */ diff --git a/api/src/main/java/org/opensearch/sql/api/spec/datetime/DatetimeOutputCastRule.java b/api/src/main/java/org/opensearch/sql/api/spec/datetime/DatetimeOutputCastRule.java index 9a7ae25e003..edc418928cb 100644 --- a/api/src/main/java/org/opensearch/sql/api/spec/datetime/DatetimeOutputCastRule.java +++ b/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. + * + *

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 fields = other.getRowType().getFieldList(); diff --git a/api/src/main/java/org/opensearch/sql/api/spec/datetime/DatetimeUdtNormalizeRule.java b/api/src/main/java/org/opensearch/sql/api/spec/datetime/DatetimeUdtNormalizeRule.java index b15d830d412..7fb8a488c6e 100644 --- a/api/src/main/java/org/opensearch/sql/api/spec/datetime/DatetimeUdtNormalizeRule.java +++ b/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. + * + *

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); diff --git a/api/src/test/java/org/opensearch/sql/api/spec/datetime/DatetimeExtensionTest.java b/api/src/test/java/org/opensearch/sql/api/spec/datetime/DatetimeExtensionTest.java index fc089150109..588d25d8d0f 100644 --- a/api/src/test/java/org/opensearch/sql/api/spec/datetime/DatetimeExtensionTest.java +++ b/api/src/test/java/org/opensearch/sql/api/spec/datetime/DatetimeExtensionTest.java @@ -172,6 +172,27 @@ public void testNonDatetimeFieldsNotWrapped() { """); } + @Test + public void testSequentialPlanCallsDoNotCorruptShuttleStack() { + // Regression test: DatetimeUdtNormalizeRule extends RelHomogeneousShuttle which inherits a + // stateful Deque 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 = diff --git a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalcitePlannerConcurrencyIT.java b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalcitePlannerConcurrencyIT.java new file mode 100644 index 00000000000..375a0f10716 --- /dev/null +++ b/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. + * + *

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} 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). + * + *

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

Run via: + * + * 

{@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
+ * }
+ */ +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 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 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 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 queries) throws Exception { + var executor = Executors.newFixedThreadPool(PARALLELISM); + try { + List> futures = new ArrayList<>(queries.size()); + AtomicInteger failures = new AtomicInteger(); + List 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 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); + } + } +}