This is an automated email from the ASF dual-hosted git repository.
yashmayya pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new c91d0e8378d Fix the shared test query server answering every query
with an error (#20096)
c91d0e8378d is described below
commit c91d0e8378d16f9b9b7fbdbdfd2560711e828368
Author: Yash Mayya <[email protected]>
AuthorDate: Thu Oct 29 15:11:48 2026 -0700
Fix the shared test query server answering every query with an error
(#20096)
---
...ionBrokerRequestHandlerFailureDetectorTest.java | 18 +++++----
.../pinot/core/transport/QueryRoutingTest.java | 33 ++++------------
.../pinot/core/transport/QueryServerTestUtils.java | 46 ++++++++++++++++------
3 files changed, 52 insertions(+), 45 deletions(-)
diff --git
a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandlerFailureDetectorTest.java
b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandlerFailureDetectorTest.java
index 7574bc25d94..3ff080b1b3e 100644
---
a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandlerFailureDetectorTest.java
+++
b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandlerFailureDetectorTest.java
@@ -19,6 +19,7 @@
package org.apache.pinot.broker.requesthandler;
import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.SettableFuture;
import java.io.IOException;
import java.net.ServerSocket;
@@ -31,6 +32,7 @@ import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Function;
import org.apache.pinot.broker.broker.AllowAllAccessControlFactory;
import org.apache.pinot.broker.queryquota.QueryQuotaManager;
import org.apache.pinot.broker.routing.manager.BrokerRoutingManager;
@@ -44,6 +46,7 @@ import org.apache.pinot.common.metrics.BrokerMetrics;
import org.apache.pinot.common.metrics.MetricValueUtils;
import org.apache.pinot.common.request.BrokerRequest;
import org.apache.pinot.core.common.datatable.DataTableBuilderFactory;
+import org.apache.pinot.core.query.request.ServerQueryRequest;
import org.apache.pinot.core.routing.ImplicitHybridTableRouteInfo;
import org.apache.pinot.core.routing.MultiClusterRoutingContext;
import org.apache.pinot.core.routing.SegmentsToQuery;
@@ -52,13 +55,13 @@ import org.apache.pinot.core.transport.QueryServer;
import org.apache.pinot.core.transport.QueryServerTestUtils;
import org.apache.pinot.core.transport.ServerInstance;
import
org.apache.pinot.core.transport.server.routing.stats.ServerRoutingStatsManager;
+import org.apache.pinot.segment.spi.memory.DataBuffer;
import org.apache.pinot.spi.accounting.ThreadAccountantUtils;
import org.apache.pinot.spi.env.PinotConfiguration;
import
org.apache.pinot.spi.eventlistener.query.BrokerQueryEventListenerFactory;
import org.apache.pinot.spi.utils.CommonConstants.Broker;
import org.apache.pinot.sql.parsers.CalciteSqlCompiler;
import org.apache.pinot.util.TestUtils;
-import org.mockito.stubbing.Answer;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.DataProvider;
@@ -192,24 +195,25 @@ public class
SingleConnectionBrokerRequestHandlerFailureDetectorTest {
/// Starts a server whose JVM is alive but whose query engine never finishes
anything -- stuck, or so busy that every
/// query times out. Its network thread still answers pings, because they do
not go through the query engine.
private ServerInstance startStuckQueryEngineServer() {
- return startQueryServer(0, invocation -> SettableFuture.<byte[]>create());
+ return startQueryServer(0, queryRequest -> SettableFuture.create());
}
/// Starts a server that answers queries straight away.
private ServerInstance startRespondingServer(int port) {
DataTable dataTable = DataTableBuilderFactory.getEmptyDataTable();
dataTable.getMetadata().put(MetadataKey.REQUEST_ID.getName(),
Long.toString(REQUEST_ID));
- byte[] responseBytes;
+ DataBuffer response;
try {
- responseBytes = dataTable.toBytes();
+ response = dataTable.toDataBuffer();
} catch (IOException e) {
throw new RuntimeException(e);
}
- return startQueryServer(port, invocation ->
Futures.immediateFuture(responseBytes));
+ return startQueryServer(port, queryRequest ->
Futures.immediateFuture(response));
}
- private ServerInstance startQueryServer(int port, Answer<?> submitAnswer) {
- QueryServer queryServer = QueryServerTestUtils.newQueryServer(port,
submitAnswer);
+ private ServerInstance startQueryServer(int port,
+ Function<ServerQueryRequest, ListenableFuture<DataBuffer>> responder) {
+ QueryServer queryServer = QueryServerTestUtils.newQueryServer(port,
responder);
_queryServers.add(queryServer);
return enable(QueryServerTestUtils.serverInstance("localhost",
QueryServerTestUtils.startAndGetPort(queryServer)));
}
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/transport/QueryRoutingTest.java
b/pinot-core/src/test/java/org/apache/pinot/core/transport/QueryRoutingTest.java
index b8516d1dc5d..755e51cbaca 100644
---
a/pinot-core/src/test/java/org/apache/pinot/core/transport/QueryRoutingTest.java
+++
b/pinot-core/src/test/java/org/apache/pinot/core/transport/QueryRoutingTest.java
@@ -20,6 +20,7 @@ package org.apache.pinot.core.transport;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.SettableFuture;
+import com.google.common.util.concurrent.Uninterruptibles;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import io.netty.channel.ChannelHandler;
@@ -31,7 +32,6 @@ import io.netty.channel.embedded.EmbeddedChannel;
import io.netty.util.ReferenceCountUtil;
import java.io.IOException;
import java.io.InputStream;
-import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.net.SocketTimeoutException;
@@ -52,12 +52,10 @@ import org.apache.pinot.common.request.InstanceRequest;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.core.common.datatable.DataTableBuilder;
import org.apache.pinot.core.common.datatable.DataTableBuilderFactory;
-import org.apache.pinot.core.query.scheduler.QueryScheduler;
import org.apache.pinot.core.routing.SegmentsToQuery;
import
org.apache.pinot.core.transport.server.routing.stats.ServerRoutingStatsManager;
import org.apache.pinot.segment.spi.memory.DataBuffer;
import org.apache.pinot.segment.spi.memory.PinotByteBuffer;
-import org.apache.pinot.server.access.AccessControl;
import org.apache.pinot.spi.accounting.ThreadAccountantUtils;
import org.apache.pinot.spi.config.table.TableType;
import org.apache.pinot.spi.env.PinotConfiguration;
@@ -74,9 +72,8 @@ import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.Test;
-import static org.mockito.ArgumentMatchers.any;
+import static
org.apache.pinot.core.transport.QueryServerTestUtils.startAndGetPort;
import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
import static org.testng.Assert.*;
@@ -149,27 +146,10 @@ public class QueryRoutingTest {
}
private QueryServer getQueryServer(int responseDelayMs, DataBuffer response,
int port) {
- InstanceRequestHandler handler = new InstanceRequestHandler("server01",
new PinotConfiguration(),
- mockQueryScheduler(responseDelayMs, response),
mock(AccessControl.class),
- ThreadAccountantUtils.getNoOpAccountant());
- return new QueryServer(port, null, handler);
- }
-
- /// Starts the query server and returns the actual bound port. Uses port 0
to let the OS assign a free port,
- /// avoiding TOCTOU race conditions. [QueryServer#start()] blocks on
`bind().sync()`, so the
- /// server is ready to accept connections when this method returns.
- private int startAndGetPort(QueryServer server) {
- server.start();
- return ((InetSocketAddress) server.getChannel().localAddress()).getPort();
- }
-
- private QueryScheduler mockQueryScheduler(int responseDelayMs, DataBuffer
response) {
- QueryScheduler queryScheduler = mock(QueryScheduler.class);
- when(queryScheduler.submitResponse(any())).thenAnswer(invocation -> {
- Thread.sleep(responseDelayMs);
+ return QueryServerTestUtils.newQueryServer(port, queryRequest -> {
+ Uninterruptibles.sleepUninterruptibly(responseDelayMs,
TimeUnit.MILLISECONDS);
return Futures.immediateFuture(response);
});
- return queryScheduler;
}
private QueryRouter newIsolatedQueryRouter() {
@@ -850,7 +830,7 @@ public class QueryRoutingTest {
/// comes back at once even while a query to that server is stuck.
@Test
public void testPingIsAnsweredPromptlyWhileTheQueryEngineIsStuck() {
- _queryServer = QueryServerTestUtils.newQueryServer(invocation ->
SettableFuture.<byte[]>create());
+ _queryServer = QueryServerTestUtils.newQueryServer(queryRequest ->
SettableFuture.create());
initializeTestFixtures(startAndGetPort(_queryServer));
QueryRouter queryRouter = newIsolatedQueryRouter();
try {
@@ -861,6 +841,9 @@ public class QueryRoutingTest {
assertTrue(queryRouter.ping(_serverInstance, 5_000L));
assertTrue(System.currentTimeMillis() - startTimeMs < 2_000L,
"A server whose query engine is stuck must still answer a ping
promptly");
+ // The status stays IN_PROGRESS until the query is waited on, even after
the server answers it, so check for an
+ // answer directly.
+ assertEquals(stuckQuery.getNumServersResponded(), 0, "The stuck query
must still be unanswered");
assertEquals(stuckQuery.getStatus(), QueryResponse.Status.IN_PROGRESS);
} finally {
queryRouter.shutDown();
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/transport/QueryServerTestUtils.java
b/pinot-core/src/test/java/org/apache/pinot/core/transport/QueryServerTestUtils.java
index 53cb513e8b6..79717ad16b7 100644
---
a/pinot-core/src/test/java/org/apache/pinot/core/transport/QueryServerTestUtils.java
+++
b/pinot-core/src/test/java/org/apache/pinot/core/transport/QueryServerTestUtils.java
@@ -18,16 +18,20 @@
*/
package org.apache.pinot.core.transport;
+import com.google.common.util.concurrent.ListenableFuture;
import java.net.InetSocketAddress;
+import java.util.concurrent.atomic.LongAccumulator;
+import java.util.function.Function;
+import org.apache.pinot.core.query.executor.QueryExecutor;
+import org.apache.pinot.core.query.request.ServerQueryRequest;
import org.apache.pinot.core.query.scheduler.QueryScheduler;
+import org.apache.pinot.core.query.scheduler.resources.ResourceManager;
+import org.apache.pinot.segment.spi.memory.DataBuffer;
import org.apache.pinot.server.access.AccessControl;
import org.apache.pinot.spi.accounting.ThreadAccountantUtils;
import org.apache.pinot.spi.env.PinotConfiguration;
-import org.mockito.stubbing.Answer;
-import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
/// Helpers for standing up a real Netty [QueryServer] in a test.
@@ -38,18 +42,34 @@ public class QueryServerTestUtils {
private QueryServerTestUtils() {
}
- /// Builds an unstarted query server whose scheduler answers every
submission with `submitAnswer`.
- public static QueryServer newQueryServer(Answer<?> submitAnswer) {
- return newQueryServer(0, submitAnswer);
+ /// Builds an unstarted query server whose scheduler answers each query with
the future `responder` returns for it.
+ ///
+ /// `responder` runs on the server's network threads, which also answer
pings. To model a slow or stuck query
+ /// engine, return a future that completes later or never: a `responder`
that blocks holds back pings as well.
+ public static QueryServer newQueryServer(Function<ServerQueryRequest,
ListenableFuture<DataBuffer>> responder) {
+ return newQueryServer(0, responder);
}
- /// Same as [#newQueryServer(Answer)], on a given port -- for a test that
brings a server back where it was.
- public static QueryServer newQueryServer(int port, Answer<?> submitAnswer) {
- QueryScheduler queryScheduler = mock(QueryScheduler.class);
- when(queryScheduler.submit(any())).thenAnswer(submitAnswer);
- InstanceRequestHandler handler =
- new InstanceRequestHandler("testServer", new PinotConfiguration(),
queryScheduler, mock(AccessControl.class),
- ThreadAccountantUtils.getNoOpAccountant());
+ /// Same as [#newQueryServer(Function)], on a given port -- for a test that
brings a server back where it was.
+ public static QueryServer newQueryServer(int port,
+ Function<ServerQueryRequest, ListenableFuture<DataBuffer>> responder) {
+ PinotConfiguration config = new PinotConfiguration();
+ // A subclass, not a mock: the scheduler sends both submit methods to the
one overridden here, so the server gets
+ // this answer whichever one it calls. A mock would answer null from the
one not stubbed.
+ QueryScheduler queryScheduler = new QueryScheduler(config, "testServer",
mock(QueryExecutor.class),
+ ThreadAccountantUtils.getNoOpAccountant(), new
LongAccumulator(Long::max, 0), mock(ResourceManager.class)) {
+ @Override
+ public ListenableFuture<DataBuffer> submitResponse(ServerQueryRequest
queryRequest) {
+ return responder.apply(queryRequest);
+ }
+
+ @Override
+ public String name() {
+ return "test";
+ }
+ };
+ InstanceRequestHandler handler = new InstanceRequestHandler("testServer",
config, queryScheduler,
+ mock(AccessControl.class), ThreadAccountantUtils.getNoOpAccountant());
return new QueryServer(port, null, handler);
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
