Allows multiple requests per server per request ID

pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java

@@ -25,6 +25,7 @@
 import java.net.URI;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -87,6 +88,7 @@
 import org.apache.pinot.core.routing.RoutingTable;
 import org.apache.pinot.core.routing.TimeBoundaryInfo;
 import org.apache.pinot.core.transport.ServerInstance;
+import org.apache.pinot.core.transport.ServerQueryRoutingContext;
 import org.apache.pinot.core.util.GapfillUtils;
 import org.apache.pinot.query.parser.utils.ParserUtils;
 import org.apache.pinot.spi.auth.AuthorizationResult;
@@ -616,44 +618,32 @@ protected BrokerResponse handleRequest(long requestId, String query, SqlNodeAndO
       // Calculate routing table for the query
       // TODO: Modify RoutingManager interface to directly take PinotQuery
       long routingStartTimeNs = System.nanoTime();
-      Map<ServerInstance, Pair<List<String>, List<String>>> offlineRoutingTable = null;
-      Map<ServerInstance, Pair<List<String>, List<String>>> realtimeRoutingTable = null;
+      Map<ServerInstance, List<ServerQueryRoutingContext>> queryRoutingTable = new HashMap<>();
       List<String> unavailableSegments = new ArrayList<>();
       int numPrunedSegmentsTotal = 0;
+
       if (offlineBrokerRequest != null) {
-        // NOTE: Routing table might be null if table is just removed
-        RoutingTable routingTable = _routingManager.getRoutingTable(offlineBrokerRequest, requestId);
-        if (routingTable != null) {
-          unavailableSegments.addAll(routingTable.getUnavailableSegments());
-          Map<ServerInstance, Pair<List<String>, List<String>>> serverInstanceToSegmentsMap =
-              routingTable.getServerInstanceToSegmentsMap();
-          if (!serverInstanceToSegmentsMap.isEmpty()) {
-            offlineRoutingTable = serverInstanceToSegmentsMap;
-          } else {
-            offlineBrokerRequest = null;
-          }
-          numPrunedSegmentsTotal += routingTable.getNumPrunedSegments();
-        } else {
+        Integer numPrunedSegments =
+            updateRoutingTable(requestId, offlineBrokerRequest, queryRoutingTable, unavailableSegments);
+        if (numPrunedSegments == null) {
           offlineBrokerRequest = null;
+        } else {
+          numPrunedSegmentsTotal += numPrunedSegments;
         }
       }
-      if (realtimeBrokerRequest != null) {
-        // NOTE: Routing table might be null if table is just removed
-        RoutingTable routingTable = _routingManager.getRoutingTable(realtimeBrokerRequest, requestId);
-        if (routingTable != null) {
-          unavailableSegments.addAll(routingTable.getUnavailableSegments());
-          Map<ServerInstance, Pair<List<String>, List<String>>> serverInstanceToSegmentsMap =
-              routingTable.getServerInstanceToSegmentsMap();
-          if (!serverInstanceToSegmentsMap.isEmpty()) {
-            realtimeRoutingTable = serverInstanceToSegmentsMap;
-          } else {
-            realtimeBrokerRequest = null;
-          }
-          numPrunedSegmentsTotal += routingTable.getNumPrunedSegments();
-        } else {
+
+      // TODO: Assess if the Explain Plan Query should also be routed to REALTIME servers for HYBRID tables
+      if (realtimeBrokerRequest != null && (!pinotQuery.isExplain() || offlineBrokerRequest != null)) {
+        // Don't send explain queries to realtime for OFFLINE or HYBRID tables
+        Integer numPrunedSegments =
+            updateRoutingTable(requestId, realtimeBrokerRequest, queryRoutingTable, unavailableSegments);
+        if (numPrunedSegments == null) {
           realtimeBrokerRequest = null;
+        } else {
+          numPrunedSegmentsTotal += numPrunedSegments;
         }
       }
+
       int numUnavailableSegments = unavailableSegments.size();
       requestContext.setNumUnavailableSegments(numUnavailableSegments);
 
@@ -716,18 +706,13 @@ protected BrokerResponse handleRequest(long requestId, String query, SqlNodeAndO
 
       // Set the maximum serialized response size per server, and ask server to directly return final response when only
       // one server is queried
-      int numServers = 0;
-      if (offlineRoutingTable != null) {
-        numServers += offlineRoutingTable.size();
-      }
-      if (realtimeRoutingTable != null) {
-        numServers += realtimeRoutingTable.size();
-      }
+      int numQueriesIssued = queryRoutingTable.values().stream().mapToInt(List::size).sum();
+
       if (offlineBrokerRequest != null) {
         Map<String, String> queryOptions = offlineBrokerRequest.getPinotQuery().getQueryOptions();
-        setMaxServerResponseSizeBytes(numServers, queryOptions, offlineTableConfig);
+        setMaxServerResponseSizeBytes(numQueriesIssued, queryOptions, offlineTableConfig);
         // Set the query option to directly return final result for single server query unless it is explicitly disabled
-        if (numServers == 1) {
+        if (numQueriesIssued == 1) {
           // Set the same flag in the original server request to be used in the reduce phase for hybrid table
           if (queryOptions.putIfAbsent(QueryOptionKey.SERVER_RETURN_FINAL_RESULT, "true") == null
               && offlineBrokerRequest != serverBrokerRequest) {
@@ -738,9 +723,9 @@ protected BrokerResponse handleRequest(long requestId, String query, SqlNodeAndO
       }
       if (realtimeBrokerRequest != null) {
         Map<String, String> queryOptions = realtimeBrokerRequest.getPinotQuery().getQueryOptions();
-        setMaxServerResponseSizeBytes(numServers, queryOptions, realtimeTableConfig);
+        setMaxServerResponseSizeBytes(numQueriesIssued, queryOptions, realtimeTableConfig);
         // Set the query option to directly return final result for single server query unless it is explicitly disabled
-        if (numServers == 1) {
+        if (numQueriesIssued == 1) {
           // Set the same flag in the original server request to be used in the reduce phase for hybrid table
           if (queryOptions.putIfAbsent(QueryOptionKey.SERVER_RETURN_FINAL_RESULT, "true") == null
               && realtimeBrokerRequest != serverBrokerRequest) {
@@ -758,15 +743,7 @@ protected BrokerResponse handleRequest(long requestId, String query, SqlNodeAndO
       //       - Compile time function invocation
       //       - Literal only queries
       //       - Any rewrites
-      if (pinotQuery.isExplain()) {
-        // Update routing tables to only send request to offline servers for OFFLINE and HYBRID tables.
-        // TODO: Assess if the Explain Plan Query should also be routed to REALTIME servers for HYBRID tables
-        if (offlineRoutingTable != null) {
-          // For OFFLINE and HYBRID tables, don't send EXPLAIN query to realtime servers.
-          realtimeBrokerRequest = null;
-          realtimeRoutingTable = null;
-        }
-      }
+
       BrokerResponseNative brokerResponse;
       if (_queriesById != null) {
         // Start to track the running query for cancellation just before sending it out to servers to avoid any
@@ -777,20 +754,20 @@ protected BrokerResponse handleRequest(long requestId, String query, SqlNodeAndO
         //       can always list the running queries and cancel query again until it ends. Just that such race
         //       condition makes cancel API less reliable. This should be rare as it assumes sending queries out to
         //       servers takes time, but will address later if needed.
-        _queriesById.put(requestId, new QueryServers(query, offlineRoutingTable, realtimeRoutingTable));
+        _queriesById.put(requestId, new QueryServers(query, queryRoutingTable));
         LOGGER.debug("Keep track of running query: {}", requestId);
         try {
-          brokerResponse = processBrokerRequest(requestId, brokerRequest, serverBrokerRequest, offlineBrokerRequest,
-              offlineRoutingTable, realtimeBrokerRequest, realtimeRoutingTable, remainingTimeMs, serverStats,
-              requestContext);
+          brokerResponse =
+              processBrokerRequest(requestId, brokerRequest, serverBrokerRequest, queryRoutingTable, remainingTimeMs,
+                  serverStats, requestContext);
         } finally {
           _queriesById.remove(requestId);
           LOGGER.debug("Remove track of running query: {}", requestId);
         }
       } else {
-        brokerResponse = processBrokerRequest(requestId, brokerRequest, serverBrokerRequest, offlineBrokerRequest,
-            offlineRoutingTable, realtimeBrokerRequest, realtimeRoutingTable, remainingTimeMs, serverStats,
-            requestContext);
+        brokerResponse =
+            processBrokerRequest(requestId, brokerRequest, serverBrokerRequest, queryRoutingTable, remainingTimeMs,
+                serverStats, requestContext);
       }
 
       for (ProcessingException exception : exceptions) {
@@ -827,6 +804,33 @@ protected BrokerResponse handleRequest(long requestId, String query, SqlNodeAndO
     }
   }
 
+  private Integer updateRoutingTable(long requestId, BrokerRequest brokerRequest,
+      Map<ServerInstance, List<ServerQueryRoutingContext>> routingTableResult, List<String> unavailableSegments) {
+    // NOTE: Routing table might be null if table is just removed
+    RoutingTable routingTable = _routingManager.getRoutingTable(brokerRequest, requestId);
+    if (routingTable != null) {
+      unavailableSegments.addAll(routingTable.getUnavailableSegments());
+      Map<ServerInstance, Pair<List<String>, List<String>>> serverInstanceToSegmentsMap =
+          routingTable.getServerInstanceToSegmentsMap();
+      if (!serverInstanceToSegmentsMap.isEmpty()) {
+        for (Map.Entry<ServerInstance, Pair<List<String>, List<String>>> serverMap
+            : serverInstanceToSegmentsMap.entrySet()) {
+          ServerInstance serverInstance = serverMap.getKey();
+          Pair<List<String>, List<String>> segmentsToQuery = serverMap.getValue();
+
+          routingTableResult.computeIfAbsent(serverInstance, k -> new ArrayList<>()).add(
+              new ServerQueryRoutingContext(brokerRequest, segmentsToQuery,
+                  serverInstance.toServerRoutingInstance(serverInstance.isTlsEnabled())));
+        }
+      } else {
+        return null;
+      }
+      return routingTable.getNumPrunedSegments();
+    } else {
+      return null;
+    }
+  }
+
   @VisibleForTesting
   static String addRoutingPolicyInErrMsg(String errorMessage, String realtimeRoutingPolicy,
       String offlineRoutingPolicy) {
@@ -903,7 +907,7 @@ private void setTimestampIndexExpressionOverrideHints(@Nullable Expression expre
       case "datetrunc":
         String granularString = function.getOperands().get(0).getLiteral().getStringValue().toUpperCase();
         Expression timeExpression = function.getOperands().get(1);
-        if (((function.getOperandsSize() == 2) || (function.getOperandsSize() == 3 && "MILLISECONDS".equalsIgnoreCase(
+        if (((function.getOperandsSize() == 2) || (function.getOperandsSize() == 3 && "MILLISECONDS" .equalsIgnoreCase(
             function.getOperands().get(2).getLiteral().getStringValue()))) && TimestampIndexUtils.isValidGranularity(
             granularString) && timeExpression.getIdentifier() != null) {
           String timeColumn = timeExpression.getIdentifier().getName();
@@ -1646,7 +1650,7 @@ private static void fixColumnName(String rawTableName, Expression expression, Ma
   @VisibleForTesting
   static String getActualColumnName(String rawTableName, String columnName, @Nullable Map<String, String> columnNameMap,
       boolean ignoreCase) {
-    if ("*".equals(columnName)) {
+    if ("*" .equals(columnName)) {
       return columnName;
     }
     String columnNameToCheck = trimTableName(rawTableName, columnName, ignoreCase);
@@ -1750,7 +1754,7 @@ private long setQueryTimeout(String tableNameWithType, Map<String, String> query
    * 5. BrokerConfig -> maxServerResponseSizeBytes
    * 6. BrokerConfig -> maxServerResponseSizeBytes
    */
-  private void setMaxServerResponseSizeBytes(int numServers, Map<String, String> queryOptions,
+  private void setMaxServerResponseSizeBytes(int numQueriesIssued, Map<String, String> queryOptions,
       @Nullable TableConfig tableConfig) {
     // QueryOption
     if (QueryOptionsUtils.getMaxServerResponseSizeBytes(queryOptions) != null) {
@@ -1759,7 +1763,7 @@ private void setMaxServerResponseSizeBytes(int numServers, Map<String, String> q
     Long maxQueryResponseSizeQueryOption = QueryOptionsUtils.getMaxQueryResponseSizeBytes(queryOptions);
     if (maxQueryResponseSizeQueryOption != null) {
       queryOptions.put(QueryOptionKey.MAX_SERVER_RESPONSE_SIZE_BYTES,
-          Long.toString(maxQueryResponseSizeQueryOption / numServers));
+          Long.toString(maxQueryResponseSizeQueryOption / numQueriesIssued));
       return;
     }
 
@@ -1773,7 +1777,7 @@ private void setMaxServerResponseSizeBytes(int numServers, Map<String, String> q
       }
       if (queryConfig.getMaxQueryResponseSizeBytes() != null) {
         queryOptions.put(QueryOptionKey.MAX_SERVER_RESPONSE_SIZE_BYTES,
-            Long.toString(queryConfig.getMaxQueryResponseSizeBytes() / numServers));
+            Long.toString(queryConfig.getMaxQueryResponseSizeBytes() / numQueriesIssued));
         return;
       }
     }
@@ -1789,7 +1793,7 @@ private void setMaxServerResponseSizeBytes(int numServers, Map<String, String> q
     String maxQueryResponseSizeBrokerConfig = _config.getProperty(Broker.CONFIG_OF_MAX_QUERY_RESPONSE_SIZE_BYTES);
     if (maxQueryResponseSizeBrokerConfig != null) {
       queryOptions.put(QueryOptionKey.MAX_SERVER_RESPONSE_SIZE_BYTES,
-          Long.toString(DataSizeUtils.toBytes(maxQueryResponseSizeBrokerConfig) / numServers));
+          Long.toString(DataSizeUtils.toBytes(maxQueryResponseSizeBrokerConfig) / numQueriesIssued));
     }
   }
 
@@ -1857,10 +1861,8 @@ private static void attachTimeBoundary(PinotQuery pinotQuery, TimeBoundaryInfo t
    * TODO: Directly take PinotQuery
    */
   protected abstract BrokerResponseNative processBrokerRequest(long requestId, BrokerRequest originalBrokerRequest,
-      BrokerRequest serverBrokerRequest, @Nullable BrokerRequest offlineBrokerRequest,
-      @Nullable Map<ServerInstance, Pair<List<String>, List<String>>> offlineRoutingTable,
-      @Nullable BrokerRequest realtimeBrokerRequest,
-      @Nullable Map<ServerInstance, Pair<List<String>, List<String>>> realtimeRoutingTable, long timeoutMs,
+      BrokerRequest serverBrokerRequest,
+      Map<ServerInstance, List<ServerQueryRoutingContext>> queryRoutingTable, long timeoutMs,
       ServerStats serverStats, RequestContext requestContext)
       throws Exception;
 
@@ -1890,14 +1892,10 @@ private static class QueryServers {
     final String _query;
     final Set<ServerInstance> _servers = new HashSet<>();
 
-    QueryServers(String query, @Nullable Map<ServerInstance, Pair<List<String>, List<String>>> offlineRoutingTable,
-        @Nullable Map<ServerInstance, Pair<List<String>, List<String>>> realtimeRoutingTable) {
+    QueryServers(String query, @Nullable Map<ServerInstance, List<ServerQueryRoutingContext>> queryRoutingTable) {
       _query = query;
-      if (offlineRoutingTable != null) {
-        _servers.addAll(offlineRoutingTable.keySet());
-      }
-      if (realtimeRoutingTable != null) {
-        _servers.addAll(realtimeRoutingTable.keySet());
+      if (queryRoutingTable != null) {
+        _servers.addAll(queryRoutingTable.keySet());
       }
     }
   }

pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BrokerRequestIdGenerator.java

@@ -19,6 +19,8 @@
 package org.apache.pinot.broker.requesthandler;
 
 import java.util.concurrent.atomic.AtomicLong;
+import org.apache.pinot.common.utils.request.BrokerRequestIdUtils;
+
 
 /**
  * An ID generator to produce a global unique identifier for each query, used in v1/v2 engine for tracking and
@@ -43,7 +45,9 @@ public BrokerRequestIdGenerator(String brokerId) {
   }
 
   public long get() {
-    long normalized = (_incrementingId.getAndIncrement() & Long.MAX_VALUE) % OFFSET;
-    return _mask + normalized;
+    long normalized =
+        ((_incrementingId.getAndIncrement() & Long.MAX_VALUE) % (OFFSET / BrokerRequestIdUtils.TABLE_TYPE_OFFSET))
+            * BrokerRequestIdUtils.TABLE_TYPE_OFFSET;
+    return BrokerRequestIdUtils.getCanonicalRequestId(_mask + normalized);
   }
 }