segmentMetadata query cancellation

processing/src/main/java/io/druid/query/metadata/SegmentMetadataQueryRunnerFactory.java

@@ -22,15 +22,21 @@ package io.druid.query.metadata;
 import com.google.common.base.Function;
 import com.google.common.base.Throwables;
 import com.google.common.collect.Maps;
-import com.metamx.common.guava.ExecutorExecutingSequence;
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.ListeningExecutorService;
+import com.google.common.util.concurrent.MoreExecutors;
+import com.google.inject.Inject;
 import com.metamx.common.guava.Sequence;
 import com.metamx.common.guava.Sequences;
+import com.metamx.common.logger.Logger;
 import io.druid.query.AbstractPrioritizedCallable;
 import io.druid.query.ConcatQueryRunner;
 import io.druid.query.Query;
+import io.druid.query.QueryInterruptedException;
 import io.druid.query.QueryRunner;
 import io.druid.query.QueryRunnerFactory;
 import io.druid.query.QueryToolChest;
+import io.druid.query.QueryWatcher;
 import io.druid.query.metadata.metadata.ColumnAnalysis;
 import io.druid.query.metadata.metadata.ColumnIncluderator;
 import io.druid.query.metadata.metadata.SegmentAnalysis;
@@ -40,16 +46,27 @@ import io.druid.segment.Segment;
 
 import java.util.Arrays;
 import java.util.Map;
-import java.util.concurrent.Callable;
+import java.util.concurrent.CancellationException;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 
 public class SegmentMetadataQueryRunnerFactory implements QueryRunnerFactory<SegmentAnalysis, SegmentMetadataQuery>
 {
   private static final SegmentAnalyzer analyzer = new SegmentAnalyzer();
-
   private static final SegmentMetadataQueryQueryToolChest toolChest = new SegmentMetadataQueryQueryToolChest();
+  private static final Logger log = new Logger(SegmentMetadataQueryRunnerFactory.class);
+
+  private final QueryWatcher queryWatcher;
+
+  @Inject
+  public SegmentMetadataQueryRunnerFactory(
+      QueryWatcher queryWatcher
+  )
+  {
+    this.queryWatcher = queryWatcher;
+  }
 
   @Override
   public QueryRunner<SegmentAnalysis> createRunner(final Segment segment)
@@ -101,9 +118,10 @@ public class SegmentMetadataQueryRunnerFactory implements QueryRunnerFactory<Seg
 
   @Override
   public QueryRunner<SegmentAnalysis> mergeRunners(
-      final ExecutorService queryExecutor, Iterable<QueryRunner<SegmentAnalysis>> queryRunners
+      ExecutorService exec, Iterable<QueryRunner<SegmentAnalysis>> queryRunners
   )
   {
+    final ListeningExecutorService queryExecutor = MoreExecutors.listeningDecorator(exec);
     return new ConcatQueryRunner<SegmentAnalysis>(
         Sequences.map(
             Sequences.simple(queryRunners),
@@ -118,7 +136,7 @@ public class SegmentMetadataQueryRunnerFactory implements QueryRunnerFactory<Seg
                   public Sequence<SegmentAnalysis> run(final Query<SegmentAnalysis> query)
                   {
                     final int priority = query.getContextPriority(0);
-                    Future<Sequence<SegmentAnalysis>> future = queryExecutor.submit(
+                    ListenableFuture<Sequence<SegmentAnalysis>> future = queryExecutor.submit(
                         new AbstractPrioritizedCallable<Sequence<SegmentAnalysis>>(priority)
                         {
                           @Override
@@ -129,13 +147,25 @@ public class SegmentMetadataQueryRunnerFactory implements QueryRunnerFactory<Seg
                         }
                     );
                     try {
-                      return future.get();
+                      queryWatcher.registerQuery(query, future);
+                      final Number timeout = query.getContextValue("timeout", (Number) null);
+                      return timeout == null ? future.get() : future.get(timeout.longValue(), TimeUnit.MILLISECONDS);
                     }
                     catch (InterruptedException e) {
-                      throw Throwables.propagate(e);
+                      log.warn(e, "Query interrupted, cancelling pending results, query id [%s]", query.getId());
+                      future.cancel(true);
+                      throw new QueryInterruptedException("Query interrupted");
+                    }
+                    catch(CancellationException e) {
+                      throw new QueryInterruptedException("Query cancelled");
+                    }
+                    catch(TimeoutException e) {
+                      log.info("Query timeout, cancelling pending results for query id [%s]", query.getId());
+                      future.cancel(true);
+                      throw new QueryInterruptedException("Query timeout");
                     }
                     catch (ExecutionException e) {
-                      throw Throwables.propagate(e);
+                      throw Throwables.propagate(e.getCause());
                     }
                   }
                 };

processing/src/test/java/io/druid/query/metadata/SegmentAnalyzerTest.java

@@ -25,6 +25,7 @@ import io.druid.query.LegacyDataSource;
 import io.druid.query.QueryRunner;
 import io.druid.query.QueryRunnerFactory;
 import io.druid.query.QueryRunnerTestHelper;
+import io.druid.query.QueryWatcher;
 import io.druid.query.metadata.metadata.ColumnAnalysis;
 import io.druid.query.metadata.metadata.SegmentAnalysis;
 import io.druid.query.metadata.metadata.SegmentMetadataQuery;
@@ -95,7 +96,7 @@ public class SegmentAnalyzerTest
   private List<SegmentAnalysis> getSegmentAnalysises(Segment index)
   {
     final QueryRunner runner = QueryRunnerTestHelper.makeQueryRunner(
-        (QueryRunnerFactory) new SegmentMetadataQueryRunnerFactory(), index
+        (QueryRunnerFactory) new SegmentMetadataQueryRunnerFactory(QueryRunnerTestHelper.NOOP_QUERYWATCHER), index
     );
 
     final SegmentMetadataQuery query = new SegmentMetadataQuery(

processing/src/test/java/io/druid/query/metadata/SegmentMetadataQueryTest.java

@@ -45,7 +45,7 @@ public class SegmentMetadataQueryTest
 {
   @SuppressWarnings("unchecked")
   private final QueryRunner runner = makeQueryRunner(
-      new SegmentMetadataQueryRunnerFactory()
+      new SegmentMetadataQueryRunnerFactory(QueryRunnerTestHelper.NOOP_QUERYWATCHER)
   );
   private ObjectMapper mapper = new DefaultObjectMapper();