Revert 'Pull Request !790 : [#I3DV3A] Move finished or failed queries to finished query states to improve HA performance'

presto-main/src/main/java/io/prestosql/dispatcher/DispatchManager.java

@@ -39,7 +39,6 @@ import io.prestosql.spi.QueryId;
 import io.prestosql.spi.resourcegroups.SelectionContext;
 import io.prestosql.spi.resourcegroups.SelectionCriteria;
 import io.prestosql.spi.service.PropertyService;
-import io.prestosql.spi.statestore.StateStore;
 import io.prestosql.statestore.SharedQueryState;
 import io.prestosql.statestore.StateCacheStore;
 import io.prestosql.statestore.StateFetcher;
@@ -65,7 +64,6 @@ import java.util.concurrent.Executor;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.Lock;
 import java.util.stream.Collectors;
-import java.util.stream.Stream;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.ImmutableList.toImmutableList;
@@ -312,11 +310,7 @@ public class DispatchManager
         List<BasicQueryInfo> queryInfos;
         if (isMultiCoordinatorEnabled() && StateCacheStore.get().getCachedStates(StateStoreConstants.QUERY_STATE_COLLECTION_NAME) != null) {
             Map<String, SharedQueryState> queryStates = StateCacheStore.get().getCachedStates(StateStoreConstants.QUERY_STATE_COLLECTION_NAME);
-            Map<String, SharedQueryState> finishedQueryStates = StateCacheStore.get().getCachedStates(StateStoreConstants.FINISHED_QUERY_STATE_COLLECTION_NAME);
-
-            queryInfos = Stream.concat(
-                    queryStates.values().stream(),
-                    finishedQueryStates.values().stream())
+            queryInfos = queryStates.values().stream()
                     .map(SharedQueryState::getBasicQueryInfo)
                     .collect(Collectors.toList());
         }
@@ -401,7 +395,6 @@ public class DispatchManager
 
         // Start state fetcher
         stateFetcher.registerStateCollection(StateStoreConstants.QUERY_STATE_COLLECTION_NAME);
-        stateFetcher.registerStateCollection(StateStoreConstants.FINISHED_QUERY_STATE_COLLECTION_NAME);
         stateFetcher.registerStateCollection(StateStoreConstants.OOM_QUERY_STATE_COLLECTION_NAME);
         stateFetcher.registerStateCollection(StateStoreConstants.CPU_USAGE_STATE_COLLECTION_NAME);
         stateFetcher.start();
@@ -423,13 +416,12 @@ public class DispatchManager
     private synchronized void submitQuerySync(DispatchQuery dispatchQuery, SelectionContext selectionContext)
             throws InterruptedException, PrestoException
     {
-        StateStore stateStore = stateStoreProvider.getStateStore();
-        if (stateStore == null) {
+        if (stateStoreProvider.getStateStore() == null) {
             LOG.error("StateStore is not loaded yet");
             throw new PrestoException(GENERIC_INTERNAL_ERROR, "Coordinator is not ready to accept queries");
         }
 
-        Lock lock = stateStore.getLock(StateStoreConstants.SUBMIT_QUERY_LOCK_NAME);
+        Lock lock = stateStoreProvider.getStateStore().getLock(StateStoreConstants.SUBMIT_QUERY_LOCK_NAME);
         // Make sure query submission is synchronized
         boolean locked = lock.tryLock(hetuConfig.getQuerySubmitTimeout().toMillis(), TimeUnit.MILLISECONDS);
         long start = 0L;
@@ -440,7 +432,7 @@ public class DispatchManager
                         dispatchQuery.getQueryId(),
                         start,
                         new SimpleDateFormat("HH:mm:ss:SSS").format(new Date(start)));
-                stateFetcher.fetchQueryStates(stateStore);
+                stateFetcher.fetchStates();
                 resourceGroupManager.submit(dispatchQuery, selectionContext, queryExecutor);
                 // Register dispatch query to StateUpdater
                 if (PropertyService.getBooleanProperty(HetuConstant.MULTI_COORDINATOR_ENABLED) && stateUpdater != null) {
@@ -449,7 +441,7 @@ public class DispatchManager
                 stateUpdater.updateStates();
             }
             catch (IOException e) {
-                throw new PrestoException(GENERIC_INTERNAL_ERROR, "Failed to fetch states from or update states to state store: " + e.getMessage());
+                throw new PrestoException(GENERIC_INTERNAL_ERROR, format("Failed to fetch states from or update states to state store: " + e.getMessage()));
             }
             catch (Throwable e) {
                 // dispatch query has already been registered, so just fail it directly
@@ -467,7 +459,7 @@ public class DispatchManager
         }
         else {
             // TODO maybe just queue the query if the queue size is not a problem
-            throw new PrestoException(GENERIC_INTERNAL_ERROR, "Coordinator probably too busy at the moment, please try again in a few minutes");
+            throw new PrestoException(GENERIC_INTERNAL_ERROR, format("Coordinator probably too busy at the moment, please try again in a few minutes"));
         }
     }
 }

presto-main/src/main/java/io/prestosql/execution/QueryTracker.java

@@ -22,7 +22,6 @@ import io.prestosql.spi.PrestoException;
 import io.prestosql.spi.QueryId;
 import io.prestosql.spi.statestore.StateCollection;
 import io.prestosql.spi.statestore.StateMap;
-import io.prestosql.spi.statestore.StateStore;
 import io.prestosql.statestore.StateStoreConstants;
 import io.prestosql.statestore.StateStoreProvider;
 import org.joda.time.DateTime;
@@ -255,11 +254,7 @@ public class QueryTracker<T extends TrackedQuery>
 
     private void removeQueryInStateStore(QueryId queryId)
     {
-        StateStore stateStore = stateStoreProvider.getStateStore();
-        if (stateStore == null) {
-            return;
-        }
-        StateCollection stateCollection = stateStore.getStateCollection(StateStoreConstants.FINISHED_QUERY_STATE_COLLECTION_NAME);
+        StateCollection stateCollection = stateStoreProvider.getStateStore().getStateCollection(StateStoreConstants.QUERY_STATE_COLLECTION_NAME);
         if (stateCollection != null && stateCollection.getType().equals(StateCollection.Type.MAP)) {
             ((StateMap<String, String>) stateCollection).remove(queryId.getId());
         }

presto-main/src/main/java/io/prestosql/statestore/LocalStateStoreProvider.java

@@ -120,7 +120,6 @@ public class LocalStateStoreProvider
         // Create essential state collections
         stateStore.createStateCollection(StateStoreConstants.DISCOVERY_SERVICE_COLLECTION_NAME, StateCollection.Type.MAP);
         stateStore.createStateCollection(StateStoreConstants.QUERY_STATE_COLLECTION_NAME, StateCollection.Type.MAP);
-        stateStore.createStateCollection(StateStoreConstants.FINISHED_QUERY_STATE_COLLECTION_NAME, StateCollection.Type.MAP);
         stateStore.createStateCollection(StateStoreConstants.OOM_QUERY_STATE_COLLECTION_NAME, StateCollection.Type.MAP);
         stateStore.createStateCollection(StateStoreConstants.CPU_USAGE_STATE_COLLECTION_NAME, StateCollection.Type.MAP);
         stateStore.createStateCollection(StateStoreConstants.TRANSACTION_STATE_COLLECTION_NAME, StateCollection.Type.MAP);

presto-main/src/main/java/io/prestosql/statestore/StateFetcher.java

@@ -24,7 +24,6 @@ import io.prestosql.server.BasicQueryInfo;
 import io.prestosql.spi.ErrorType;
 import io.prestosql.spi.statestore.StateCollection;
 import io.prestosql.spi.statestore.StateMap;
-import io.prestosql.spi.statestore.StateStore;
 import org.joda.time.DateTime;
 import org.joda.time.DateTimeZone;
 
@@ -44,12 +43,6 @@ import java.util.concurrent.locks.Lock;
 import static com.google.common.base.Preconditions.checkState;
 import static io.airlift.concurrent.Threads.threadsNamed;
 import static io.prestosql.spi.StandardErrorCode.SERVER_SHUTTING_DOWN;
-import static io.prestosql.statestore.StateStoreConstants.CPU_USAGE_STATE_COLLECTION_NAME;
-import static io.prestosql.statestore.StateStoreConstants.DEFAULT_ACQUIRED_LOCK_TIME_MS;
-import static io.prestosql.statestore.StateStoreConstants.FINISHED_QUERY_STATE_COLLECTION_NAME;
-import static io.prestosql.statestore.StateStoreConstants.HANDLE_EXPIRED_QUERY_LOCK_NAME;
-import static io.prestosql.statestore.StateStoreConstants.OOM_QUERY_STATE_COLLECTION_NAME;
-import static io.prestosql.statestore.StateStoreConstants.QUERY_STATE_COLLECTION_NAME;
 import static io.prestosql.utils.StateUtils.removeState;
 
 /**
@@ -154,7 +147,7 @@ public class StateFetcher
                 if (stateCollection == null) {
                     continue;
                 }
-                if (stateCollectionName.equals(CPU_USAGE_STATE_COLLECTION_NAME)) {
+                if (stateCollectionName.equals(StateStoreConstants.CPU_USAGE_STATE_COLLECTION_NAME)) {
                     StateCacheStore.get().setCachedStates(stateCollectionName, ((StateMap) stateCollection).getAll());
                     continue;
                 }
@@ -176,46 +169,6 @@ public class StateFetcher
                     LOG.warn("Unsupported state collection type: %s", stateCollection.getType());
                 }
             }
-            long end = System.currentTimeMillis();
-            LOG.debug("fetchStates ends at current time milliseconds: %s, at format HH:mm:ss:SSS:%s, total time use: %s",
-                    end,
-                    new SimpleDateFormat("HH:mm:ss:SSS").format(new Date(end)),
-                    end - start);
-        }
-    }
-
-    /**
-     * Fetch state from state store to cache store
-     *
-     * @throws IOException exception when failed to deserialize states
-     */
-    public void fetchQueryStates(StateStore stateStore)
-            throws IOException
-    {
-        synchronized (this) {
-            long start = System.currentTimeMillis();
-            LOG.debug("fetchStates starts at current time milliseconds: %s, at format HH:mm:ss:SSS:%s",
-                    start,
-                    new SimpleDateFormat("HH:mm:ss:SSS").format(new Date(start)));
-
-            DateTime currentTime = new DateTime(DateTimeZone.UTC);
-            StateCollection cpuUsageCollection = stateStore.getStateCollection(CPU_USAGE_STATE_COLLECTION_NAME);
-            StateCollection queryStateCollection = stateStore.getStateCollection(QUERY_STATE_COLLECTION_NAME);
-
-            StateCacheStore.get().setCachedStates(CPU_USAGE_STATE_COLLECTION_NAME, ((StateMap) cpuUsageCollection).getAll());
-
-            Map<String, String> states = ((StateMap<String, String>) queryStateCollection).getAll();
-
-            ImmutableMap.Builder<String, SharedQueryState> queryStatesBuilder = ImmutableMap.builder();
-            for (Map.Entry<String, String> entry : states.entrySet()) {
-                SharedQueryState state = MAPPER.readerFor(SharedQueryState.class).readValue(entry.getValue());
-                if (isStateExpired(state, currentTime)) {
-                    handleExpiredQueryState(state);
-                }
-                queryStatesBuilder.put(entry.getKey(), state);
-            }
-            StateCacheStore.get().setCachedStates(QUERY_STATE_COLLECTION_NAME, queryStatesBuilder.build());
-
             long end = System.currentTimeMillis();
             LOG.debug("updateStates ends at current time milliseconds: %s, at format HH:mm:ss:SSS:%s, total time use: %s",
                     end,
@@ -244,16 +197,15 @@ public class StateFetcher
     private void handleExpiredQueryState(SharedQueryState state)
     {
         // State store hasn't been loaded yet
-        final StateStore stateStore = stateStoreProvider.getStateStore();
-        if (stateStore == null) {
+        if (stateStoreProvider.getStateStore() == null) {
             return;
         }
 
         Lock lock = null;
         boolean locked = false;
         try {
-            lock = stateStore.getLock(HANDLE_EXPIRED_QUERY_LOCK_NAME);
-            locked = lock.tryLock(DEFAULT_ACQUIRED_LOCK_TIME_MS, TimeUnit.MILLISECONDS);
+            lock = stateStoreProvider.getStateStore().getLock(StateStoreConstants.HANDLE_EXPIRED_QUERY_LOCK_NAME);
+            locked = lock.tryLock(StateStoreConstants.DEFAULT_ACQUIRED_LOCK_TIME_MS, TimeUnit.MILLISECONDS);
             if (locked) {
                 LOG.debug(String.format("EXPIRED!!! REMOVING... Id: %s, state: %s, uri: %s, query: %s",
                         state.getBasicQueryInfo().getQueryId().getId(),
@@ -262,11 +214,11 @@ public class StateFetcher
                         state.getBasicQueryInfo().getQuery()));
 
                 // remove expired query from oom
-                StateCollection stateCollection = stateStore.getStateCollection(OOM_QUERY_STATE_COLLECTION_NAME);
+                StateCollection stateCollection = stateStoreProvider.getStateStore().getStateCollection(StateStoreConstants.OOM_QUERY_STATE_COLLECTION_NAME);
                 removeState(stateCollection, Optional.of(state.getBasicQueryInfo().getQueryId()), LOG);
 
                 // update query to failed in stateCollection if exists
-                stateCollection = stateStore.getStateCollection(FINISHED_QUERY_STATE_COLLECTION_NAME);
+                stateCollection = stateStoreProvider.getStateStore().getStateCollection(StateStoreConstants.QUERY_STATE_COLLECTION_NAME);
                 if (stateCollection != null && stateCollection.getType().equals(StateCollection.Type.MAP)) {
                     Map<String, String> queryStateMap = ((StateMap<String, String>) stateCollection).getAll();
                     if (queryStateMap.get(state.getBasicQueryInfo().getQueryId().getId()) != null) {

presto-main/src/main/java/io/prestosql/statestore/StateStoreConstants.java

@@ -50,11 +50,6 @@ public class StateStoreConstants
      */
     public static final String QUERY_STATE_COLLECTION_NAME = "query";
 
-    /**
-     * Finished query state collection name
-     */
-    public static final String FINISHED_QUERY_STATE_COLLECTION_NAME = "finished-query";
-
     /**
      * OOM Query state collection name
      */

presto-main/src/main/java/io/prestosql/statestore/StateUpdater.java

@@ -18,22 +18,22 @@ import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.collect.ArrayListMultimap;
 import com.google.common.collect.Multimap;
-import com.google.common.collect.Multimaps;
 import io.airlift.json.ObjectMapperProvider;
 import io.airlift.log.Logger;
 import io.airlift.units.Duration;
 import io.prestosql.dispatcher.DispatchQuery;
 import io.prestosql.execution.ManagedQueryExecution;
 import io.prestosql.execution.QueryState;
+import io.prestosql.spi.QueryId;
 import io.prestosql.spi.statestore.StateCollection;
 import io.prestosql.spi.statestore.StateMap;
-import io.prestosql.spi.statestore.StateStore;
 
 import java.text.SimpleDateFormat;
 import java.util.Date;
-import java.util.LinkedList;
-import java.util.List;
+import java.util.HashSet;
+import java.util.Iterator;
 import java.util.Optional;
+import java.util.Set;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
@@ -43,9 +43,6 @@ import static com.google.common.base.Preconditions.checkState;
 import static io.airlift.concurrent.Threads.threadsNamed;
 import static io.prestosql.spi.StandardErrorCode.CLUSTER_OUT_OF_MEMORY;
 import static io.prestosql.spi.StandardErrorCode.EXCEEDED_GLOBAL_MEMORY_LIMIT;
-import static io.prestosql.statestore.StateStoreConstants.FINISHED_QUERY_STATE_COLLECTION_NAME;
-import static io.prestosql.statestore.StateStoreConstants.OOM_QUERY_STATE_COLLECTION_NAME;
-import static io.prestosql.statestore.StateStoreConstants.QUERY_STATE_COLLECTION_NAME;
 import static io.prestosql.utils.StateUtils.removeState;
 
 /**
@@ -59,7 +56,7 @@ public class StateUpdater
 
     private final StateStoreProvider stateStoreProvider;
     private final Duration updateInterval;
-    private final Multimap<String, DispatchQuery> registeredQueries = Multimaps.synchronizedMultimap(ArrayListMultimap.create());
+    private final Multimap<String, DispatchQuery> states = ArrayListMultimap.create();
     private final ScheduledExecutorService stateUpdateExecutor;
     private ScheduledFuture<?> backgroundTask;
 
@@ -97,7 +94,7 @@ public class StateUpdater
         synchronized (this) {
             if (backgroundTask != null) {
                 backgroundTask.cancel(true);
-                registeredQueries.clear();
+                states.clear();
             }
         }
     }
@@ -111,14 +108,14 @@ public class StateUpdater
      */
     public void registerQuery(String stateCollectionName, DispatchQuery query)
     {
-        synchronized (registeredQueries) {
-            registeredQueries.put(stateCollectionName, query);
+        synchronized (states) {
+            states.put(stateCollectionName, query);
+            query.addStateChangeListener(state -> {
+                if (state.isDone()) {
+                    queryFinished(query);
+                }
+            });
         }
-        query.addStateChangeListener(state -> {
-            if (state.isDone()) {
-                queryFinished(query);
-            }
-        });
     }
 
     /**
@@ -129,8 +126,8 @@ public class StateUpdater
      */
     public void unregisterQuery(String stateCollectionName, ManagedQueryExecution query)
     {
-        synchronized (registeredQueries) {
-            registeredQueries.remove(stateCollectionName, query);
+        synchronized (states) {
+            states.remove(stateCollectionName, query);
         }
     }
 
@@ -142,55 +139,54 @@ public class StateUpdater
     public void updateStates()
             throws JsonProcessingException
     {
-        // State store hasn't been loaded yet
-        final StateStore stateStore = stateStoreProvider.getStateStore();
-        if (stateStore == null) {
-            return;
-        }
+        synchronized (states) {
+            // State store hasn't been loaded yet
+            if (stateStoreProvider.getStateStore() == null) {
+                return;
+            }
 
-        long start = System.currentTimeMillis();
-        LOG.debug("UpdateStates starts at current time milliseconds: %s, at format HH:mm:ss:SSS:%s",
-                start,
-                new SimpleDateFormat("HH:mm:ss:SSS").format(new Date(start)));
+            long start = System.currentTimeMillis();
+            LOG.debug("UpdateStates starts at current time milliseconds: %s, at format HH:mm:ss:SSS:%s",
+                    start,
+                    new SimpleDateFormat("HH:mm:ss:SSS").format(new Date(start)));
 
-        StateCollection finishedQueries = stateStore.getStateCollection(FINISHED_QUERY_STATE_COLLECTION_NAME);
-        StateCollection queries = stateStore.getStateCollection(QUERY_STATE_COLLECTION_NAME);
+            for (String stateCollectionName : states.keySet()) {
+                StateCollection stateCollection = stateStoreProvider.getStateStore().getStateCollection(stateCollectionName);
+                Set<QueryId> finishedQueries = new HashSet<>();
+                for (DispatchQuery query : states.get(stateCollectionName)) {
+                    SharedQueryState state = SharedQueryState.create(query);
+                    String stateJson = MAPPER.writeValueAsString(state);
 
-        List<DispatchQuery> queriesToUnregister = new LinkedList<>();
-        synchronized (registeredQueries) {
-            for (DispatchQuery query : registeredQueries.get(QUERY_STATE_COLLECTION_NAME)) {
-                SharedQueryState state = SharedQueryState.create(query);
-                String stateJson = MAPPER.writeValueAsString(state);
+                    switch (stateCollection.getType()) {
+                        case MAP:
+                            ((StateMap) stateCollection).put(state.getBasicQueryInfo().getQueryId().getId(), stateJson);
+                            break;
+                        default:
+                            LOG.error("Unsupported state collection type: %s", stateCollection.getType());
+                    }
 
-                if (state.getBasicQueryInfo().getState() == QueryState.FINISHED || state.getBasicQueryInfo().getState() == QueryState.FAILED) {
-                    // No need to update states for finished queries
-                    // also move finished queries to finished-query state collection
-                    queriesToUnregister.add(query);
-                    ((StateMap) finishedQueries).put(state.getBasicQueryInfo().getQueryId().getId(), stateJson);
-                    continue;
+                    if (state.getBasicQueryInfo().getState() == QueryState.FINISHED || state.getBasicQueryInfo().getState() == QueryState.FAILED) {
+                        finishedQueries.add(state.getBasicQueryInfo().getQueryId());
+                    }
                 }
 
-                ((StateMap) queries).put(state.getBasicQueryInfo().getQueryId().getId(), stateJson);
+                // No need to update states for finished queries
+                unregisterFinishedQueries(stateCollectionName, finishedQueries);
             }
-        }
 
-        for (DispatchQuery query : queriesToUnregister) {
-            removeFromStateCollection(stateStore, QUERY_STATE_COLLECTION_NAME, query);
+            long end = System.currentTimeMillis();
+            LOG.debug("updateStates ends at current time milliseconds: %s, at format HH:mm:ss:SSS:%s, total time use: %s",
+                    end,
+                    new SimpleDateFormat("HH:mm:ss:SSS").format(new Date(end)),
+                    end - start);
         }
-
-        long end = System.currentTimeMillis();
-        LOG.debug("updateStates ends at current time milliseconds: %s, at format HH:mm:ss:SSS:%s, total time use: %s",
-                end,
-                new SimpleDateFormat("HH:mm:ss:SSS").format(new Date(end)),
-                end - start);
     }
 
     private void queryFinished(ManagedQueryExecution query)
     {
-        StateStore stateStore = stateStoreProvider.getStateStore();
         // If query killed by OOM remove the query from OOM query state store
-        if (stateStore != null && isQueryKilledByOOMKiller(query)) {
-            removeFromStateCollection(stateStore, OOM_QUERY_STATE_COLLECTION_NAME, query);
+        if (isQueryKilledByOOMKiller(query)) {
+            removeFromStateStore(StateStoreConstants.OOM_QUERY_STATE_COLLECTION_NAME, query);
         }
     }
 
@@ -200,14 +196,42 @@ public class StateUpdater
             return false;
         }
 
-        return query.getErrorCode().get().equals(CLUSTER_OUT_OF_MEMORY.toErrorCode()) ||
-                query.getErrorCode().get().equals(EXCEEDED_GLOBAL_MEMORY_LIMIT.toErrorCode());
+        if (query.getErrorCode().get().equals(CLUSTER_OUT_OF_MEMORY.toErrorCode()) ||
+                query.getErrorCode().get().equals(EXCEEDED_GLOBAL_MEMORY_LIMIT.toErrorCode())) {
+            return true;
+        }
+
+        return false;
     }
 
-    private void removeFromStateCollection(StateStore stateStore, String stateCollectionName, ManagedQueryExecution query)
+    private void removeFromStateStore(String stateCollectionName, ManagedQueryExecution query)
     {
-        StateCollection stateCollection = stateStore.getStateCollection(stateCollectionName);
-        removeState(stateCollection, Optional.of(query.getBasicQueryInfo().getQueryId()), LOG);
-        unregisterQuery(stateCollectionName, query);
+        // State store hasn't been loaded yet
+        if (stateStoreProvider.getStateStore() == null) {
+            return;
+        }
+
+        synchronized (states) {
+            StateCollection stateCollection = stateStoreProvider.getStateStore().getStateCollection(stateCollectionName);
+            removeState(stateCollection, Optional.of(query.getBasicQueryInfo().getQueryId()), LOG);
+            states.remove(stateCollectionName, query);
+        }
+    }
+
+    /**
+     * For finished or failed queries, no need to keep updating their states to state store
+     * so unregister them from state updater
+     *
+     * @param stateCollectionName state collection name
+     * @param queryIds Queries to be unregistered
+     */
+    private void unregisterFinishedQueries(String stateCollectionName, Set<QueryId> queryIds)
+    {
+        Iterator<DispatchQuery> iterator = states.get(stateCollectionName).iterator();
+        while (iterator.hasNext()) {
+            if (queryIds.contains(iterator.next().getQueryId())) {
+                iterator.remove();
+            }
+        }
     }
 }

presto-main/src/main/java/io/prestosql/utils/DistributedResourceGroupUtils.java

@@ -86,7 +86,7 @@ public class DistributedResourceGroupUtils
      */
     public static void accumulateCpuUsage(StateStore stateStore)
     {
-        StateCollection queryStateCollection = stateStore.getStateCollection(StateStoreConstants.FINISHED_QUERY_STATE_COLLECTION_NAME);
+        StateCollection queryStateCollection = stateStore.getStateCollection(StateStoreConstants.QUERY_STATE_COLLECTION_NAME);
         if (queryStateCollection == null) {
             return;
         }

presto-main/src/test/java/io/prestosql/statestore/TestStateUpdater.java

@@ -68,6 +68,7 @@ public class TestStateUpdater
 {
     private static final int MINIMUM_UPDATE_INTERVAL = 10;
     private static final String STATE_COLLECTION_QUERY = "query";
+    private static final String DISPATCH_QUERY_ENTRY = "dispatchquery";
     private static final String NULL_STRING = null;
     private static final String MOCK_QUERY_ID = "20191122_174317_00000_q9aun";
     private static final String URI_LOCALHOST = "http://localhost:8080";
@@ -222,7 +223,7 @@ public class TestStateUpdater
         StateUpdater stateUpdater = new StateUpdater(stateStoreProvider, updateInterval);
         DispatchQuery dispatchQuery = mockDispatchQueryData(true);
         updateStateChange(dispatchQuery);
-        stateUpdater.registerQuery(STATE_COLLECTION_QUERY, dispatchQuery);
+        stateUpdater.registerQuery(DISPATCH_QUERY_ENTRY, dispatchQuery);
         when(stateStoreProvider.getStateStore()).then(new Returns(stateStore));
         when(stateStoreProvider.getStateStore().getStateCollection(any())).then(new Returns(Mockito.mock(StateMap.class)));
         when(stateStoreProvider.getStateStore().getStateCollection(any()).getType()).then(new Returns(StateCollection.Type.MAP));

presto-main/src/test/java/io/prestosql/utils/TestDistributedResourceGroupUtils.java

@@ -152,7 +152,7 @@ public class TestDistributedResourceGroupUtils
 
             ResourceGroupId root = new ResourceGroupId("root");
 
-            StateStore stateStore = setupMockStateStore(new HashMap<>(), new HashMap<>(), new HashMap<>());
+            StateStore stateStore = setupMockStateStore(new HashMap<>(), new HashMap<>());
 
             //query1 completed, query2 completed, query3 running
             MockManagedQueryExecution query1 = new MockManagedQueryExecution(0, "query1", 0,
@@ -178,8 +178,8 @@ public class TestDistributedResourceGroupUtils
             SharedQueryState queryState2 = getSharedQueryState(query2);
             SharedQueryState queryState3 = getSharedQueryState(query3);
 
-            ((StateMap) stateStore.getStateCollection(StateStoreConstants.FINISHED_QUERY_STATE_COLLECTION_NAME)).put(query1.toString(), mapper.writeValueAsString(queryState1));
-            ((StateMap) stateStore.getStateCollection(StateStoreConstants.FINISHED_QUERY_STATE_COLLECTION_NAME)).put(query2.toString(), mapper.writeValueAsString(queryState2));
+            ((StateMap) stateStore.getStateCollection(StateStoreConstants.QUERY_STATE_COLLECTION_NAME)).put(query1.toString(), mapper.writeValueAsString(queryState1));
+            ((StateMap) stateStore.getStateCollection(StateStoreConstants.QUERY_STATE_COLLECTION_NAME)).put(query2.toString(), mapper.writeValueAsString(queryState2));
             ((StateMap) stateStore.getStateCollection(StateStoreConstants.QUERY_STATE_COLLECTION_NAME)).put(query3.toString(), mapper.writeValueAsString(queryState3));
 
             DistributedResourceGroupUtils.accumulateCpuUsage(stateStore);
@@ -196,28 +196,21 @@ public class TestDistributedResourceGroupUtils
         }
     }
 
-    private static StateStore setupMockStateStore(Map<String, String> queryMap, Map<String, String> finishedQueryMap, Map<String, Long> cpuUsageMap)
+    private static StateStore setupMockStateStore(Map<String, String> queryMap, Map<String, Long> cpuUsageMap)
     {
         StateMap mockQueryMap = mock(StateMap.class);
-        StateMap mockFinishedQueryMap = mock(StateMap.class);
         StateMap mockCpuUsageMap = mock(StateMap.class);
 
         when(mockQueryMap.put(anyString(), anyString())).thenAnswer(i -> queryMap.put((String) i.getArguments()[0], (String) i.getArguments()[1]));
-        when(mockFinishedQueryMap.put(anyString(), anyString())).thenAnswer(i -> finishedQueryMap.put((String) i.getArguments()[0], (String) i.getArguments()[1]));
         when(mockCpuUsageMap.put(anyString(), anyString())).thenAnswer(i -> cpuUsageMap.put((String) i.getArguments()[0], (Long) i.getArguments()[1]));
 
         when(mockQueryMap.get(anyString())).thenAnswer(i -> queryMap.get(i.getArguments()[0]));
         when(mockQueryMap.getAll()).thenReturn(queryMap);
-
-        when(mockFinishedQueryMap.get(anyString())).thenAnswer(i -> finishedQueryMap.get(i.getArguments()[0]));
-        when(mockFinishedQueryMap.getAll()).thenReturn(finishedQueryMap);
-
         when(mockCpuUsageMap.get(anyString())).thenAnswer(i -> cpuUsageMap.get(i.getArguments()[0]));
         when(mockCpuUsageMap.getAll()).thenReturn(cpuUsageMap);
 
         StateStore stateStore = mock(StateStore.class);
         when(stateStore.getStateCollection(StateStoreConstants.QUERY_STATE_COLLECTION_NAME)).thenReturn(mockQueryMap);
-        when(stateStore.getStateCollection(StateStoreConstants.FINISHED_QUERY_STATE_COLLECTION_NAME)).thenReturn(mockFinishedQueryMap);
         when(stateStore.getStateCollection(StateStoreConstants.CPU_USAGE_STATE_COLLECTION_NAME)).thenReturn(mockCpuUsageMap);
         return stateStore;
     }