public boolean cancelQuery(String queryId) {
    Preconditions.checkState(_queryFuturesById != null, "Query cancellation is not enabled on server");
    // Keep the future as it'll be cleaned up by the thread executing the query.
    Future<byte[]> future = _queryFuturesById.get(queryId);
    if (future == null) {
      return false;
    }
    boolean done = future.isDone();
    if (!done) {
      future.cancel(true);
    }
    if (LOGGER.isDebugEnabled()) {
      LOGGER.debug("Cancelled query: {} that's done: {}", queryId, done);
    }
    return true;
  }
@Test
  public void testCancelQuery() {
    PinotConfiguration config = new PinotConfiguration();
    config.setProperty("pinot.server.enable.query.cancellation", "true");
    QueryScheduler qs = createQueryScheduler(config);
    InstanceRequestHandler handler =
        new InstanceRequestHandler("server01", config, qs, mock(ServerMetrics.class), mock(AccessControl.class));

    Set<String> queryIds = new HashSet<>();
    queryIds.add("foo");
    queryIds.add("bar");
    queryIds.add("baz");
    for (String id : queryIds) {
      ServerQueryRequest query = mock(ServerQueryRequest.class);
      when(query.getQueryId()).thenReturn(id);
      ChannelHandlerContext ctx = mock(ChannelHandlerContext.class);
      ChannelFuture chFu = mock(ChannelFuture.class);
      when(ctx.writeAndFlush(any())).thenReturn(chFu);
      handler.submitQuery(query, ctx, "myTable01", System.currentTimeMillis(), mock(InstanceRequest.class));
    }
    Assert.assertEquals(handler.getRunningQueryIds(), queryIds);
    for (String id : queryIds) {
      handler.cancelQuery(id);
    }
    Assert.assertTrue(handler.getRunningQueryIds().isEmpty());
    Assert.assertFalse(handler.cancelQuery("unknown"));
  }