yashmayya commented on code in PR #19867: URL: https://github.com/apache/pinot/pull/19867#discussion_r4275142644
########## pinot-core/src/test/java/org/apache/pinot/core/operator/combine/CombineTerminationTest.java: ########## @@ -0,0 +1,236 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.core.operator.combine; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.function.Function; +import org.apache.pinot.core.common.Block; +import org.apache.pinot.core.common.Operator; +import org.apache.pinot.core.operator.BaseOperator; +import org.apache.pinot.core.operator.ExecutionStatistics; +import org.apache.pinot.core.operator.blocks.results.BaseResultsBlock; +import org.apache.pinot.core.operator.streaming.BaseStreamingCombineOperator; +import org.apache.pinot.core.operator.streaming.StreamingSelectionOnlyCombineOperator; +import org.apache.pinot.core.query.request.context.QueryContext; +import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils; +import org.apache.pinot.spi.accounting.ThreadAccountantUtils; +import org.apache.pinot.spi.exception.QueryErrorCode; +import org.apache.pinot.spi.exception.QueryErrorMessage; +import org.apache.pinot.spi.query.QueryExecutionContext; +import org.apache.pinot.spi.query.QueryThreadContext; +import org.apache.pinot.util.TestUtils; +import org.testng.annotations.AfterMethod; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertTrue; + + +/// Tests that a combine operator returns promptly when its query is terminated, even when the worker tasks are +/// cancelled while still queued in the executor. A cancelled queued task never runs, so it never signals completion +/// to the main thread. The main thread must observe the termination itself instead of waiting until the query +/// deadline while it keeps the merged results reachable. +@SuppressWarnings("rawtypes") +public class CombineTerminationTest { + private static final int NUM_OPERATORS = 10; + private static final int NUM_THREADS = 2; + private static final long QUERY_TIMEOUT_MS = 60_000L; + private static final long MAX_RETURN_TIME_MS = 10_000L; + // The combine operator checks termination every 100 ms. Leave a margin for slow CI machines. + private static final long MAX_TERMINATION_LATENCY_MS = 5_000L; + + private ThreadPoolExecutor _workerPool; + private CountDownLatch _releaseBlockers; + private ExecutorService _combineExecutor; + + @BeforeMethod + public void setUp() + throws InterruptedException { + // Saturate the worker pool, so that every task the combine operator submits stays queued + _workerPool = (ThreadPoolExecutor) Executors.newFixedThreadPool(NUM_THREADS); + _releaseBlockers = new CountDownLatch(1); + CountDownLatch blockersStarted = new CountDownLatch(NUM_THREADS); + for (int i = 0; i < NUM_THREADS; i++) { + _workerPool.submit(() -> { + blockersStarted.countDown(); + _releaseBlockers.await(); + return null; + }); + } + assertTrue(blockersStarted.await(MAX_RETURN_TIME_MS, TimeUnit.MILLISECONDS)); + // Avoid early finalization by not using Executors.newSingleThreadExecutor (java <= 20, JDK-8145304) + _combineExecutor = Executors.newFixedThreadPool(1); + } + + @AfterMethod + public void tearDown() { + _releaseBlockers.countDown(); + _workerPool.shutdownNow(); + _combineExecutor.shutdownNow(); + } + + @Test + public void testSelectionOnlyCombineOperator() + throws Exception { + testTerminateWithQueuedTasks("SELECT * FROM testTable", + (args) -> new SelectionOnlyCombineOperator(args._operators, args._queryContext, args._executorService)); + } + + @Test + public void testSelectionOrderByCombineOperator() + throws Exception { + testTerminateWithQueuedTasks("SELECT * FROM testTable ORDER BY column", + (args) -> new SelectionOrderByCombineOperator(args._operators, args._queryContext, args._executorService)); + } + + @Test + public void testAggregationCombineOperator() + throws Exception { + testTerminateWithQueuedTasks("SELECT COUNT(*) FROM testTable", + (args) -> new AggregationCombineOperator(args._operators, args._queryContext, args._executorService)); + } + + @Test + public void testGroupByCombineOperator() + throws Exception { + testTerminateWithQueuedTasks("SELECT COUNT(*) FROM testTable GROUP BY column", + (args) -> new GroupByCombineOperator(args._operators, args._queryContext, args._executorService)); + } + + @Test + public void testSortedGroupByCombineOperator() + throws Exception { + testTerminateWithQueuedTasks( + "SET sortAggregateSingleThreadedNumSegmentsThreshold=1; " + + "SELECT column, COUNT(*) FROM testTable GROUP BY column ORDER BY column LIMIT 10", + (args) -> new SortedGroupByCombineOperator(args._operators, args._queryContext, args._executorService)); + } + + @Test + public void testSequentialSortedGroupByCombineOperator() + throws Exception { + testTerminateWithQueuedTasks("SELECT column, COUNT(*) FROM testTable GROUP BY column ORDER BY column LIMIT 10", + (args) -> new SequentialSortedGroupByCombineOperator(args._operators, args._queryContext, + args._executorService)); + } + + @Test + public void testStreamingSelectionOnlyCombineOperator() + throws Exception { + testTerminateWithQueuedTasks("SELECT * FROM testTable", + (args) -> new StreamingSelectionOnlyCombineOperator(args._operators, args._queryContext, + args._executorService)); + } + + private void testTerminateWithQueuedTasks(String query, Function<CombineArgs, BaseCombineOperator<?>> factory) + throws Exception { + List<Operator> operators = new ArrayList<>(NUM_OPERATORS); + for (int i = 0; i < NUM_OPERATORS; i++) { + operators.add(new NeverScheduledOperator()); + } + QueryContext queryContext = QueryContextConverterUtils.getQueryContext(query); + queryContext.setEndTimeMs(System.currentTimeMillis() + QUERY_TIMEOUT_MS); + QueryExecutionContext executionContext = QueryExecutionContext.forSseTest(); + CountDownLatch combineStarted = new CountDownLatch(1); + + // Run the combine operator like the query scheduler does: on a thread with its own query context, and with a + // context aware executor that registers the worker tasks for cancellation + Future<BaseResultsBlock> future = _combineExecutor.submit(() -> { + try (QueryThreadContext ignore = QueryThreadContext.open(executionContext, + ThreadAccountantUtils.getNoOpAccountant())) { + ExecutorService executorService = QueryThreadContext.contextAwareExecutorService(_workerPool); + BaseCombineOperator<?> combineOperator = factory.apply(new CombineArgs(operators, queryContext, + executorService)); + combineStarted.countDown(); + if (combineOperator instanceof BaseStreamingCombineOperator) { + ((BaseStreamingCombineOperator<?>) combineOperator).start(); + } + return (BaseResultsBlock) combineOperator.nextBlock(); + } + }); + assertTrue(combineStarted.await(MAX_RETURN_TIME_MS, TimeUnit.MILLISECONDS)); + // Wait until the combine operator has submitted its tasks, so that termination cancels them while they are queued + TestUtils.waitForCondition(aVoid -> !_workerPool.getQueue().isEmpty(), MAX_RETURN_TIME_MS, Review Comment: Nit (not blocking): this wait ends when the first task is in the queue. At that time, the main thread can still be in `startProcess()`. Then `terminate()` runs before the wait starts, and the 100 ms loop is not tested. What I saw locally: - Without the fix, 6 of 7 cases fail, but `testStreamingSelectionOnlyCombineOperator` passes. `BaseStreamingCombineOperator` already calls `checkTermination()` before the poll. - With a fix that checks only once, before the loop, 3 of 7 cases pass. Suggest: keep the main thread in a field, and wait until its state is `TIMED_WAITING` before you call `terminate()`. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
