From def62c74f8e60f799b7118d26554e5f340ca9cc5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Xavier=20L=C3=A9aut=C3=A9?= <xavier@metamarkets.com>
Date: Wed, 14 May 2014 17:01:18 -0700
Subject: [PATCH] properly remove completed queries from query manager

---
 .../java/io/druid/server/QueryManager.java    | 52 +++++++++++++------
 1 file changed, 36 insertions(+), 16 deletions(-)

diff --git a/server/src/main/java/io/druid/server/QueryManager.java b/server/src/main/java/io/druid/server/QueryManager.java
index 5910aeda535..84b947414dd 100644
--- a/server/src/main/java/io/druid/server/QueryManager.java
+++ b/server/src/main/java/io/druid/server/QueryManager.java
@@ -19,30 +19,50 @@
 
 package io.druid.server;
 
-import com.google.common.collect.Maps;
+import com.google.common.collect.HashMultimap;
+import com.google.common.collect.Multimaps;
+import com.google.common.collect.SetMultimap;
 import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.MoreExecutors;
 import io.druid.query.Query;
 import io.druid.query.QueryWatcher;
 
-import java.util.concurrent.ConcurrentMap;
-import java.util.concurrent.Future;
+import java.util.Set;
 
 public class QueryManager implements QueryWatcher
 {
-  final ConcurrentMap<String, ListenableFuture> queries;
+  final SetMultimap<String, ListenableFuture> queries;
 
-  public QueryManager() {
-    this.queries = Maps.newConcurrentMap();
-  }
-
-  public void cancelQuery(String id) {
-    Future future = queries.get(id);
-    if(future != null) {
-      future.cancel(true);
-    }
-  }
-  public void registerQuery(Query query, ListenableFuture future)
+  public QueryManager()
   {
-    queries.put(query.getId(), future);
+    this.queries = Multimaps.synchronizedSetMultimap(
+        HashMultimap.<String, ListenableFuture>create()
+    );
+  }
+
+  public boolean cancelQuery(String id) {
+    Set<ListenableFuture> futures = queries.removeAll(id);
+    boolean success = true;
+    for (ListenableFuture future : futures) {
+      success = success && future.cancel(true);
+    }
+    return success;
+  }
+
+  public void registerQuery(Query query, final ListenableFuture future)
+  {
+    final String id = query.getId();
+    queries.put(id, future);
+    future.addListener(
+        new Runnable()
+        {
+          @Override
+          public void run()
+          {
+            queries.remove(id, future);
+          }
+        },
+        MoreExecutors.sameThreadExecutor()
+    );
   }
 }