@VisibleForTesting
  void submit(long requestId, DispatchableSubPlan dispatchableSubPlan, long timeoutMs, Map<String, String> queryOptions)
      throws Exception {
    Deadline deadline = Deadline.after(timeoutMs, TimeUnit.MILLISECONDS);

    // Serialize the stage plans in parallel
    List<DispatchablePlanFragment> stagePlans = dispatchableSubPlan.getQueryStageList();
    Set<QueryServerInstance> serverInstances = new HashSet<>();
    // Ignore the reduce stage (stage 0)
    int numStages = stagePlans.size() - 1;
    List<CompletableFuture<StageInfo>> stageInfoFutures = new ArrayList<>(numStages);
    for (int i = 0; i < numStages; i++) {
      DispatchablePlanFragment stagePlan = stagePlans.get(i + 1);
      serverInstances.addAll(stagePlan.getServerInstanceToWorkerIdMap().keySet());
      stageInfoFutures.add(CompletableFuture.supplyAsync(() -> {
        ByteString rootNode = PlanNodeSerializer.process(stagePlan.getPlanFragment().getFragmentRoot()).toByteString();
        ByteString customProperty = QueryPlanSerDeUtils.toProtoProperties(stagePlan.getCustomProperties());
        return new StageInfo(rootNode, customProperty);
      }, _executorService));
    }
    List<StageInfo> stageInfos = new ArrayList<>(numStages);
    try {
      for (CompletableFuture<StageInfo> future : stageInfoFutures) {
        stageInfos.add(future.get(deadline.timeRemaining(TimeUnit.MILLISECONDS), TimeUnit.MILLISECONDS));
      }
    } finally {
      for (CompletableFuture<?> future : stageInfoFutures) {
        if (!future.isDone()) {
          future.cancel(true);
        }
      }
    }

    Map<String, String> requestMetadata = new HashMap<>();
    requestMetadata.put(CommonConstants.Query.Request.MetadataKeys.REQUEST_ID, Long.toString(requestId));
    requestMetadata.put(CommonConstants.Broker.Request.QueryOptionKey.TIMEOUT_MS,
        Long.toString(deadline.timeRemaining(TimeUnit.MILLISECONDS)));
    requestMetadata.putAll(queryOptions);
    ByteString protoRequestMetadata = QueryPlanSerDeUtils.toProtoProperties(requestMetadata);

    // Submit the query plan to all servers in parallel
    int numServers = serverInstances.size();
    BlockingQueue<AsyncQueryDispatchResponse> dispatchCallbacks = new ArrayBlockingQueue<>(numServers);
    for (QueryServerInstance serverInstance : serverInstances) {
      _executorService.submit(() -> {
        try {
          Worker.QueryRequest.Builder requestBuilder = Worker.QueryRequest.newBuilder();
          requestBuilder.setVersion(CommonConstants.MultiStageQueryRunner.PlanVersions.V1);
          for (int i = 0; i < numStages; i++) {
            int stageId = i + 1;
            DispatchablePlanFragment stagePlan = stagePlans.get(stageId);
            List<Integer> workerIds = stagePlan.getServerInstanceToWorkerIdMap().get(serverInstance);
            if (workerIds != null) {
              List<WorkerMetadata> stageWorkerMetadataList = stagePlan.getWorkerMetadataList();
              List<WorkerMetadata> workerMetadataList = new ArrayList<>(workerIds.size());
              for (int workerId : workerIds) {
                workerMetadataList.add(stageWorkerMetadataList.get(workerId));
              }
              List<Worker.WorkerMetadata> protoWorkerMetadataList =
                  QueryPlanSerDeUtils.toProtoWorkerMetadataList(workerMetadataList);
              StageInfo stageInfo = stageInfos.get(i);
              Worker.StageMetadata stageMetadata =
                  Worker.StageMetadata.newBuilder().setStageId(stageId).addAllWorkerMetadata(protoWorkerMetadataList)
                      .setCustomProperty(stageInfo._customProperty).build();
              requestBuilder.addStagePlan(
                  Worker.StagePlan.newBuilder().setRootNode(stageInfo._rootNode).setStageMetadata(stageMetadata)
                      .build());
            }
          }
          requestBuilder.setMetadata(protoRequestMetadata);
          getOrCreateDispatchClient(serverInstance).submit(requestBuilder.build(), serverInstance, deadline,
              dispatchCallbacks::offer);
        } catch (Throwable t) {
          LOGGER.warn("Caught exception while dispatching query: {} to server: {}", requestId, serverInstance, t);
          dispatchCallbacks.offer(new AsyncQueryDispatchResponse(serverInstance, null, t));
        }
      });
    }

    int numSuccessCalls = 0;
    // TODO: Cancel all dispatched requests if one of the dispatch errors out or deadline is breached.
    while (!deadline.isExpired() && numSuccessCalls < numServers) {
      AsyncQueryDispatchResponse resp =
          dispatchCallbacks.poll(deadline.timeRemaining(TimeUnit.MILLISECONDS), TimeUnit.MILLISECONDS);
      if (resp != null) {
        if (resp.getThrowable() != null) {
          throw new RuntimeException(
              String.format("Error dispatching query: %d to server: %s", requestId, resp.getServerInstance()),
              resp.getThrowable());
        } else {
          Worker.QueryResponse response = resp.getQueryResponse();
          assert response != null;
          if (response.containsMetadata(CommonConstants.Query.Response.ServerResponseStatus.STATUS_ERROR)) {
            throw new RuntimeException(
                String.format("Unable to execute query plan for request: %d on server: %s, ERROR: %s", requestId,
                    resp.getServerInstance(),
                    response.getMetadataOrDefault(CommonConstants.Query.Response.ServerResponseStatus.STATUS_ERROR,
                        "null")));
          }
          numSuccessCalls++;
        }
      }
    }
    if (deadline.isExpired()) {
      throw new TimeoutException("Timed out waiting for response of async query-dispatch");
    }
  }
@Test
  public void testQueryDispatcherThrowsWhenQueryServerCallsOnError() {
    String sql = "SELECT * FROM a WHERE col1 = 'foo'";
    QueryServer failingQueryServer = _queryServerMap.values().iterator().next();
    Mockito.doAnswer(invocationOnMock -> {
      StreamObserver<Worker.QueryResponse> observer = invocationOnMock.getArgument(1);
      observer.onError(new RuntimeException("foo"));
      return null;
    }).when(failingQueryServer).submit(Mockito.any(), Mockito.any());
    DispatchableSubPlan dispatchableSubPlan = _queryEnvironment.planQuery(sql);
    try {
      _queryDispatcher.submit(REQUEST_ID_GEN.getAndIncrement(), dispatchableSubPlan, 10_000L, Collections.emptyMap());
      Assert.fail("Method call above should have failed");
    } catch (Exception e) {
      Assert.assertTrue(e.getMessage().contains("Error dispatching query"));
    }
    Mockito.reset(failingQueryServer);
  }