diff --git a/presto-main/src/main/java/com/facebook/presto/cost/HistoryBasedPlanStatisticsCalculator.java b/presto-main/src/main/java/com/facebook/presto/cost/HistoryBasedPlanStatisticsCalculator.java
index 6e186d570a1f..a79b108b849c 100644
--- a/presto-main/src/main/java/com/facebook/presto/cost/HistoryBasedPlanStatisticsCalculator.java
+++ b/presto-main/src/main/java/com/facebook/presto/cost/HistoryBasedPlanStatisticsCalculator.java
@@ -79,10 +79,15 @@ public PlanNodeStatsEstimate calculateStats(PlanNode node, StatsProvider sourceS
     @Override
     public boolean registerPlan(PlanNode root, Session session, long startTimeInNano, long timeoutInMilliseconds)
     {
+        // If previous registration timeout for this query, this run is likely to timeout too, return false.
+        if (historyBasedStatisticsCacheManager.historyBasedQueryRegistrationTimeout(session.getQueryId())) {
+            return false;
+        }
         ImmutableList.Builder<PlanNodeWithHash> planNodesWithHash = ImmutableList.builder();
         Iterable<PlanNode> planNodeIterable = forTree(PlanNode::getSources).depthFirstPreOrder(root);
         for (PlanNode plan : planNodeIterable) {
             if (checkTimeOut(startTimeInNano, timeoutInMilliseconds)) {
+                historyBasedStatisticsCacheManager.setHistoryBasedQueryRegistrationTimeout(session.getQueryId());
                 return false;
             }
             if (plan.getStatsEquivalentPlanNode().isPresent()) {
@@ -95,6 +100,10 @@ public boolean registerPlan(PlanNode root, Session session, long startTimeInNano
         catch (ExecutionException e) {
             throw new RuntimeException("Unable to register plan: ", e.getCause());
         }
+
+        if (checkTimeOut(startTimeInNano, timeoutInMilliseconds)) {
+            historyBasedStatisticsCacheManager.setHistoryBasedQueryRegistrationTimeout(session.getQueryId());
+        }
         // Return true even if get empty history statistics, so that HistoricalStatisticsEquivalentPlanMarkingOptimizer still return the plan with StatsEquivalentPlanNode which
         // will be used in populating history statistics
         return true;
@@ -142,7 +151,7 @@ private Map<PlanCanonicalizationStrategy, PlanNodeWithHash> getPlanNodeHashes(Pl
 
     private PlanNodeStatsEstimate getStatistics(PlanNode planNode, Session session, Lookup lookup, PlanNodeStatsEstimate delegateStats)
     {
-        if (!useHistoryBasedPlanStatisticsEnabled(session) || historyBasedStatisticsCacheManager.loadHistoryFailed(session.getQueryId())) {
+        if (!useHistoryBasedPlanStatisticsEnabled(session)) {
             return delegateStats;
         }
 
diff --git a/presto-main/src/main/java/com/facebook/presto/cost/HistoryBasedStatisticsCacheManager.java b/presto-main/src/main/java/com/facebook/presto/cost/HistoryBasedStatisticsCacheManager.java
index d0c8dbf43b4b..88b288bcbfa3 100644
--- a/presto-main/src/main/java/com/facebook/presto/cost/HistoryBasedStatisticsCacheManager.java
+++ b/presto-main/src/main/java/com/facebook/presto/cost/HistoryBasedStatisticsCacheManager.java
@@ -46,8 +46,8 @@ public class HistoryBasedStatisticsCacheManager
 
     private final Map<QueryId, Map<CachingPlanCanonicalInfoProvider.InputTableCacheKey, PlanStatistics>> inputTableStatistics = new ConcurrentHashMap<>();
 
-    // Whether the history is loaded successfully
-    private final Set<QueryId> invalidQueryId = ConcurrentHashMap.newKeySet();
+    // Stores query IDs which timeout during history optimizer registration
+    private final Set<QueryId> queryIdsRegistrationTimeOut = ConcurrentHashMap.newKeySet();
 
     public HistoryBasedStatisticsCacheManager() {}
 
@@ -59,24 +59,13 @@ public LoadingCache<PlanNodeWithHash, HistoricalPlanStatistics> getStatisticsCac
                     @Override
                     public HistoricalPlanStatistics load(PlanNodeWithHash key)
                     {
-                        if (invalidQueryId.contains(queryId)) {
-                            return empty();
-                        }
                         return loadAll(Collections.singleton(key)).values().stream().findAny().orElseGet(HistoricalPlanStatistics::empty);
                     }
 
                     @Override
                     public Map<PlanNodeWithHash, HistoricalPlanStatistics> loadAll(Iterable<? extends PlanNodeWithHash> keys)
                     {
-                        if (invalidQueryId.contains(queryId)) {
-                            Map<PlanNodeWithHash, HistoricalPlanStatistics> emptyResult = new HashMap<>();
-                            keys.forEach(x -> emptyResult.put(x, empty()));
-                            return emptyResult;
-                        }
                         Map<PlanNodeWithHash, HistoricalPlanStatistics> statistics = new HashMap<>(historyBasedPlanStatisticsProvider.get().getStats(ImmutableList.copyOf(keys), timeoutInMilliSeconds));
-                        if (statistics.isEmpty()) {
-                            invalidQueryId.add(queryId);
-                        }
                         // loadAll excepts all keys to be written
                         for (PlanNodeWithHash key : keys) {
                             statistics.putIfAbsent(key, empty());
@@ -101,7 +90,7 @@ public void invalidate(QueryId queryId)
         statisticsCache.remove(queryId);
         canonicalInfoCache.remove(queryId);
         inputTableStatistics.remove(queryId);
-        invalidQueryId.remove(queryId);
+        queryIdsRegistrationTimeOut.remove(queryId);
     }
 
     @VisibleForTesting
@@ -110,8 +99,13 @@ public Map<QueryId, Map<CachingPlanCanonicalInfoProvider.CacheKey, PlanNodeCanon
         return canonicalInfoCache;
     }
 
-    public boolean loadHistoryFailed(QueryId queryId)
+    public boolean historyBasedQueryRegistrationTimeout(QueryId queryId)
+    {
+        return queryIdsRegistrationTimeOut.contains(queryId);
+    }
+
+    public void setHistoryBasedQueryRegistrationTimeout(QueryId queryId)
     {
-        return invalidQueryId.contains(queryId);
+        queryIdsRegistrationTimeOut.add(queryId);
     }
 }