Skip to content

Introduce QueryExecutionContext to manage query life cycle #16728

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Jackie-Jiang merged 1 commit into from
Oct 1, 2025
Merged

Introduce QueryExecutionContext to manage query life cycle #16728

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 6 additions & 6 deletions pinot-broker/src/main/java/org/apache/pinot/broker/api/resources/PinotBrokerDebug.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,10 +61,9 @@
import org.apache.pinot.core.transport.server.routing.stats.ServerRoutingStatsEntry;
import org.apache.pinot.core.transport.server.routing.stats.ServerRoutingStatsManager;
import org.apache.pinot.spi.accounting.QueryResourceTracker;
import org.apache.pinot.spi.accounting.ThreadAccountant;
import org.apache.pinot.spi.accounting.ThreadResourceTracker;
import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant;
import org.apache.pinot.spi.config.table.TableType;
import org.apache.pinot.spi.trace.Tracing;
import org.apache.pinot.spi.utils.builder.TableNameBuilder;
import org.apache.pinot.sql.parsers.CalciteSqlCompiler;

Expand Down Expand Up @@ -100,6 +99,9 @@ public class PinotBrokerDebug {
@Inject
AccessControlFactory _accessControlFactory;

@Inject
ThreadAccountant _threadAccountant;

@GET
@Produces(MediaType.APPLICATION_JSON)
@Path("/debug/timeBoundary/{tableName}")
Expand Down Expand Up @@ -290,8 +292,7 @@ private long getRequestId() {
@Authorize(targetType = TargetType.CLUSTER, action = Actions.Cluster.DEBUG_RESOURCE_USAGE)
@ApiOperation(value = "Get resource usage of threads")
public Collection<? extends ThreadResourceTracker> getThreadResourceUsage() {
ThreadResourceUsageAccountant threadAccountant = Tracing.getThreadAccountant();
return threadAccountant.getThreadResources();
return _threadAccountant.getThreadResources();
}

@GET
Expand All @@ -301,8 +302,7 @@ public Collection<? extends ThreadResourceTracker> getThreadResourceUsage() {
@ApiOperation(value = "Get current resource usage of queries in this service", notes = "This is a debug endpoint, "
+ "and won't maintain backward compatibility")
public Collection<? extends QueryResourceTracker> getQueryUsage() {
ThreadResourceUsageAccountant threadAccountant = Tracing.getThreadAccountant();
return threadAccountant.getQueryResources().values();
return _threadAccountant.getQueryResources().values();
}

@GET
Expand Down
6 changes: 1 addition & 5 deletions pinot-broker/src/main/java/org/apache/pinot/broker/api/resources/PinotClientRequest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,7 +83,6 @@
import org.apache.pinot.spi.env.PinotConfiguration;
import org.apache.pinot.spi.exception.QueryErrorCode;
import org.apache.pinot.spi.exception.QueryException;
import org.apache.pinot.spi.query.QueryThreadContext;
import org.apache.pinot.spi.trace.RequestContext;
import org.apache.pinot.spi.trace.RequestScope;
import org.apache.pinot.spi.trace.Tracing;
Expand Down Expand Up @@ -452,11 +451,9 @@ public String cancelQuery(
@DefaultValue("3000") int timeoutMs,
@ApiParam(value = "Return server responses for troubleshooting") @QueryParam("verbose") @DefaultValue("false")
boolean verbose) {
try (QueryThreadContext.CloseableContext closeMe = QueryThreadContext.open(_instanceId)) {
try {
Map<String, Integer> serverResponses = verbose ? new HashMap<>() : null;
if (isClient) {
long reqId = _requestHandler.getRequestIdByClientId(id).orElse(-1L);
QueryThreadContext.setIds(reqId, id);
Copy link
Contributor

@gortiz gortiz Sep 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a con of this refactor. With the previous model we had the ability to partially create the thread local object. This means that any log printed after this line would include the CID, which is may be very useful.

It is not a critial problem, but something that would be great if we find a way to keep

Copy link
Contributor Author

@Jackie-Jiang Jackie-Jiang Sep 15, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't follow why we need this to be thread local. It should be shared across all threads for a given query

if (_requestHandler.cancelQueryByClientId(id, timeoutMs, _executor, _httpConnMgr, serverResponses)) {
String resp = "Cancelled client query: " + id;
if (verbose) {
Expand All @@ -467,7 +464,6 @@ public String cancelQuery(
} else {
long reqId = Long.parseLong(id);
if (_requestHandler.cancelQuery(reqId, timeoutMs, _executor, _httpConnMgr, serverResponses)) {
QueryThreadContext.setIds(reqId, id);
String resp = "Cancelled query: " + id;
if (verbose) {
resp += " with responses from servers: " + serverResponses;
Expand Down
5 changes: 4 additions & 1 deletion pinot-broker/src/main/java/org/apache/pinot/broker/broker/BrokerAdminApiApplication.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@
import org.apache.pinot.core.transport.ListenerConfig;
import org.apache.pinot.core.transport.server.routing.stats.ServerRoutingStatsManager;
import org.apache.pinot.core.util.ListenerConfigUtil;
import org.apache.pinot.spi.accounting.ThreadAccountant;
import org.apache.pinot.spi.env.PinotConfiguration;
import org.apache.pinot.spi.utils.CommonConstants;
import org.apache.pinot.spi.utils.PinotReflectionUtils;
Expand Down Expand Up @@ -77,7 +78,8 @@ public class BrokerAdminApiApplication extends ResourceConfig {
public BrokerAdminApiApplication(BrokerRoutingManager routingManager, BrokerRequestHandler brokerRequestHandler,
BrokerMetrics brokerMetrics, PinotConfiguration brokerConf, SqlQueryExecutor sqlQueryExecutor,
ServerRoutingStatsManager serverRoutingStatsManager, AccessControlFactory accessFactory,
HelixManager helixManager, QueryQuotaManager queryQuotaManager, AbstractResponseStore responseStore) {
HelixManager helixManager, QueryQuotaManager queryQuotaManager, ThreadAccountant threadAccountant,
AbstractResponseStore responseStore) {
_brokerResourcePackages = brokerConf.getProperty(CommonConstants.Broker.BROKER_RESOURCE_PACKAGES,
CommonConstants.Broker.DEFAULT_BROKER_RESOURCE_PACKAGES);
String[] pkgs = _brokerResourcePackages.split(",");
Expand Down Expand Up @@ -118,6 +120,7 @@ protected void configure() {
bind(queryQuotaManager).to(QueryQuotaManager.class);
bind(accessFactory).to(AccessControlFactory.class);
bind(startTime).named(BrokerAdminApiApplication.START_TIME);
bind(threadAccountant).to(ThreadAccountant.class);
bind(responseStore).to(AbstractResponseStore.class);
bind(brokerConf).to(PinotConfiguration.class);
}
Expand Down
38 changes: 18 additions & 20 deletions pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,8 @@
import org.apache.pinot.query.mailbox.MailboxService;
import org.apache.pinot.query.service.dispatch.QueryDispatcher;
import org.apache.pinot.segment.local.function.GroovyFunctionEvaluator;
import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant;
import org.apache.pinot.spi.accounting.ThreadAccountant;
import org.apache.pinot.spi.accounting.ThreadAccountantUtils;
import org.apache.pinot.spi.accounting.ThreadResourceUsageProvider;
import org.apache.pinot.spi.accounting.WorkloadBudgetManager;
import org.apache.pinot.spi.cursors.ResponseStoreService;
Expand All @@ -99,7 +100,6 @@
import org.apache.pinot.spi.metrics.PinotMetricsRegistry;
import org.apache.pinot.spi.services.ServiceRole;
import org.apache.pinot.spi.services.ServiceStartable;
import org.apache.pinot.spi.trace.Tracing;
import org.apache.pinot.spi.utils.CommonConstants;
import org.apache.pinot.spi.utils.CommonConstants.Broker;
import org.apache.pinot.spi.utils.CommonConstants.Helix;
Expand Down Expand Up @@ -164,7 +164,7 @@ public abstract class BaseBrokerStarter implements ServiceStartable {
protected AbstractResponseStore _responseStore;
protected BrokerGrpcServer _brokerGrpcServer;
protected FailureDetector _failureDetector;
protected ThreadResourceUsageAccountant _resourceUsageAccountant;
protected ThreadAccountant _threadAccountant;

@Override
public void init(PinotConfiguration brokerConf)
Expand Down Expand Up @@ -365,14 +365,13 @@ public void start()
ThreadResourceUsageProvider.setThreadMemoryMeasurementEnabled(
_brokerConf.getProperty(CommonConstants.Broker.CONFIG_OF_ENABLE_THREAD_ALLOCATED_BYTES_MEASUREMENT,
CommonConstants.Broker.DEFAULT_THREAD_ALLOCATED_BYTES_MEASUREMENT));
// initialize the thread accountant for query killing
PinotConfiguration threadAccountantConfigs = _brokerConf.subset(CommonConstants.PINOT_QUERY_SCHEDULER_PREFIX);
// This allows for custom implementations of WorkloadBudgetManager.
WorkloadBudgetManager workloadBudgetManager = createWorkloadBudgetManager(threadAccountantConfigs);
_resourceUsageAccountant = Tracing.ThreadAccountantOps.createThreadAccountant(
threadAccountantConfigs, _instanceId,
org.apache.pinot.spi.config.instance.InstanceType.BROKER, workloadBudgetManager);
Preconditions.checkNotNull(_resourceUsageAccountant);
// Initialize workload budget manager and thread resource usage accountant. Workload budget manager must be
// initialized first because it might be used by the accountant.
PinotConfiguration schedulerConfig = _brokerConf.subset(CommonConstants.PINOT_QUERY_SCHEDULER_PREFIX);
WorkloadBudgetManager.set(createWorkloadBudgetManager(schedulerConfig));
_threadAccountant = ThreadAccountantUtils.createAccountant(schedulerConfig, _instanceId,
org.apache.pinot.spi.config.instance.InstanceType.BROKER);
_threadAccountant.startWatcherTask();

// Create Broker request handler.
String brokerId = _brokerConf.getProperty(Broker.CONFIG_OF_BROKER_ID, getDefaultBrokerId());
Expand All @@ -383,7 +382,7 @@ public void start()
if (brokerRequestHandlerType.equalsIgnoreCase(Broker.GRPC_BROKER_REQUEST_HANDLER_TYPE)) {
singleStageBrokerRequestHandler =
new GrpcBrokerRequestHandler(_brokerConf, brokerId, requestIdGenerator, _routingManager,
_accessControlFactory, _queryQuotaManager, tableCache, _failureDetector, _resourceUsageAccountant);
_accessControlFactory, _queryQuotaManager, tableCache, _failureDetector, _threadAccountant);
} else {
// Default request handler type, i.e. netty
NettyConfig nettyDefaults = NettyConfig.extractNettyConfig(_brokerConf, Broker.BROKER_NETTY_PREFIX);
Expand All @@ -395,7 +394,7 @@ public void start()
singleStageBrokerRequestHandler =
new SingleConnectionBrokerRequestHandler(_brokerConf, brokerId, requestIdGenerator, _routingManager,
_accessControlFactory, _queryQuotaManager, tableCache, nettyDefaults, tlsDefaults,
_serverRoutingStatsManager, _failureDetector, _resourceUsageAccountant);
_serverRoutingStatsManager, _failureDetector, _threadAccountant);
}
MultiStageBrokerRequestHandler multiStageBrokerRequestHandler = null;
QueryDispatcher queryDispatcher = null;
Expand All @@ -409,14 +408,14 @@ public void start()
multiStageBrokerRequestHandler =
new MultiStageBrokerRequestHandler(_brokerConf, brokerId, requestIdGenerator, _routingManager,
_accessControlFactory, _queryQuotaManager, tableCache, _multiStageQueryThrottler, _failureDetector,
_resourceUsageAccountant);
_threadAccountant);
}
TimeSeriesRequestHandler timeSeriesRequestHandler = null;
if (StringUtils.isNotBlank(_brokerConf.getProperty(PinotTimeSeriesConfiguration.getEnabledLanguagesConfigKey()))) {
Preconditions.checkNotNull(queryDispatcher, "Multistage Engine should be enabled to use time-series engine");
timeSeriesRequestHandler =
new TimeSeriesRequestHandler(_brokerConf, brokerId, requestIdGenerator, _routingManager,
_accessControlFactory, _queryQuotaManager, tableCache, queryDispatcher, _resourceUsageAccountant);
_accessControlFactory, _queryQuotaManager, tableCache, queryDispatcher, _threadAccountant);
}

LOGGER.info("Initializing PinotFSFactory");
Expand All @@ -440,8 +439,6 @@ public void start()
timeSeriesRequestHandler, _responseStore);
_brokerRequestHandler.start();

Tracing.ThreadAccountantOps.startThreadAccountant();

String controllerUrl = _brokerConf.getProperty(Broker.CONTROLLER_URL);
if (controllerUrl != null) {
_sqlQueryExecutor = new SqlQueryExecutor(controllerUrl);
Expand Down Expand Up @@ -721,6 +718,7 @@ public void stop() {

LOGGER.info("Shutting down request handler and broker admin application");
_brokerRequestHandler.shutDown();
_threadAccountant.stopWatcherTask();
_brokerAdminApplication.stop();

LOGGER.info("Stopping the broker routing manager");
Expand Down Expand Up @@ -779,9 +777,9 @@ protected BrokerAdminApiApplication createBrokerAdminApp() {
BrokerAdminApiApplication brokerAdminApiApplication =
new BrokerAdminApiApplication(_routingManager, _brokerRequestHandler, _brokerMetrics, _brokerConf,
_sqlQueryExecutor, _serverRoutingStatsManager, _accessControlFactory, _spectatorHelixManager,
_queryQuotaManager, _responseStore);
brokerAdminApiApplication.register(new AuditServiceBinder(_defaultClusterConfigChangeHandler, getServiceRole(),
_brokerMetrics));
_queryQuotaManager, _threadAccountant, _responseStore);
brokerAdminApiApplication.register(
new AuditServiceBinder(_defaultClusterConfigChangeHandler, getServiceRole(), _brokerMetrics));
registerExtraComponents(brokerAdminApiApplication);
return brokerAdminApiApplication;
}
Expand Down
9 changes: 5 additions & 4 deletions ...ain/java/org/apache/pinot/broker/broker/helix/BrokerUserDefinedMessageHandlerFactory.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,8 @@
import org.apache.pinot.common.messages.SegmentRefreshMessage;
import org.apache.pinot.common.messages.TableConfigRefreshMessage;
import org.apache.pinot.common.utils.DatabaseUtils;
import org.apache.pinot.spi.accounting.WorkloadBudgetManager;
import org.apache.pinot.spi.config.workload.InstanceCost;
import org.apache.pinot.spi.trace.Tracing;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

Expand Down Expand Up @@ -284,14 +284,15 @@ public HelixTaskResult handleMessage() {
LOGGER.info("Handling query workload message: {}", _message);
try {
if (_messageType.equals(QueryWorkloadRefreshMessage.DELETE_QUERY_WORKLOAD_MSG_SUB_TYPE)) {
Tracing.ThreadAccountantOps.getWorkloadBudgetManager().deleteWorkload(_queryWorkloadName);
WorkloadBudgetManager.get().deleteWorkload(_queryWorkloadName);
} else if (_messageType.equals(QueryWorkloadRefreshMessage.REFRESH_QUERY_WORKLOAD_MSG_SUB_TYPE)) {
if (_instanceCost == null) {
throw new IllegalStateException(
"Instance cost is not provided for refreshing query workload: " + _queryWorkloadName);
}
Tracing.ThreadAccountantOps.getWorkloadBudgetManager()
.addOrUpdateWorkload(_queryWorkloadName, _instanceCost.getCpuCostNs(), _instanceCost.getMemoryCostBytes());
WorkloadBudgetManager.get()
.addOrUpdateWorkload(_queryWorkloadName, _instanceCost.getCpuCostNs(),
_instanceCost.getMemoryCostBytes());
} else {
throw new IllegalStateException("Unknown message type: " + _messageType);
}
Expand Down
Loading
Loading