public static <T> void runTasksWithDeadline(int numTasks, Function<Integer, T> taskFunc, Consumer<T> resCollector,
      Consumer<Exception> errHandler, ExecutorService executorService, long deadlineInMs) {
    Phaser phaser = new Phaser(1);
    List<Future<T>> futures = new ArrayList<>(numTasks);
    for (int i = 0; i < numTasks; i++) {
      int index = i;
      futures.add(executorService.submit(new TraceCallable<T>() {
        @Override
        public T callJob() {
          try {
            // Register the thread to the phaser for main thread to wait for it to complete.
            if (phaser.register() < 0) {
              return null;
            }
            return taskFunc.apply(index);
          } finally {
            phaser.arriveAndDeregister();
          }
        }
      }));
    }
    try {
      // Check deadline while waiting for the task results.
      for (Future<T> future : futures) {
        T taskResult = future.get(deadlineInMs - System.currentTimeMillis(), TimeUnit.MILLISECONDS);
        resCollector.accept(taskResult);
      }
    } catch (Exception e) {
      errHandler.accept(e);
    } finally {
      // Cancel all ongoing jobs
      for (Future<T> future : futures) {
        if (!future.isDone()) {
          future.cancel(true);
        }
      }
      // Deregister the main thread and wait for all threads done
      phaser.awaitAdvance(phaser.arriveAndDeregister());
    }
  }
@Test
  public void testRunTasksWithDeadline() {
    ExecutorService exec = Executors.newCachedThreadPool();
    AtomicInteger sum = new AtomicInteger(0);
    QueryMultiThreadingUtils.runTasksWithDeadline(5, index -> index, sum::addAndGet, e -> {
    }, exec, System.currentTimeMillis() + 500);
    // sum of 0, 1, .., 4 indices.
    Assert.assertEquals(sum.get(), 10);

    // Task throws exception before timeout.
    Exception[] err = new Exception[1];
    QueryMultiThreadingUtils.runTasksWithDeadline(5, index -> {
      throw new RuntimeException("oops: " + index);
    }, sum::addAndGet, e -> err[0] = e, exec, System.currentTimeMillis() + 500);
    Assert.assertTrue(err[0].getMessage().contains("oops"));

    // Task timed out.
    QueryMultiThreadingUtils.runTasksWithDeadline(5, index -> {
      try {
        Thread.sleep(10_000);
        return index;
      } catch (InterruptedException e) {
        throw new RuntimeException(e);
      }
    }, sum::addAndGet, e -> err[0] = e, exec, System.currentTimeMillis() + 500);
    Assert.assertTrue(err[0] instanceof TimeoutException);
  }