Skip HBO registration if previous registration timeout

Currently we run HBO optimizer twice. When the first run timeout,
it's likely that the second run will timeout too. Here we record
the timeout, and skip registration for the next run.

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;

presto-main/src/main/java/com/facebook/presto/cost/HistoryBasedStatisticsCacheManager.java

@@ -46,6 +46,9 @@ public class HistoryBasedStatisticsCacheManager
 
     private final Map<QueryId, Map<CachingPlanCanonicalInfoProvider.InputTableCacheKey, PlanStatistics>> inputTableStatistics = new ConcurrentHashMap<>();
 
+    // Stores query IDs which timeout during history optimizer registration
+    private final Set<QueryId> queryIdsRegistrationTimeOut = ConcurrentHashMap.newKeySet();
+
     public HistoryBasedStatisticsCacheManager() {}
 
     public LoadingCache<PlanNodeWithHash, HistoricalPlanStatistics> getStatisticsCache(QueryId queryId, Supplier<HistoryBasedPlanStatisticsProvider> historyBasedPlanStatisticsProvider, long timeoutInMilliSeconds)
@@ -87,11 +90,22 @@ public void invalidate(QueryId queryId)
         statisticsCache.remove(queryId);
         canonicalInfoCache.remove(queryId);
         inputTableStatistics.remove(queryId);
+        queryIdsRegistrationTimeOut.remove(queryId);
     }
 
     @VisibleForTesting
     public Map<QueryId, Map<CachingPlanCanonicalInfoProvider.CacheKey, PlanNodeCanonicalInfo>> getCanonicalInfoCache()
     {
         return canonicalInfoCache;
     }
+
+    public boolean historyBasedQueryRegistrationTimeout(QueryId queryId)
+    {
+        return queryIdsRegistrationTimeOut.contains(queryId);
+    }
+
+    public void setHistoryBasedQueryRegistrationTimeout(QueryId queryId)
+    {
+        queryIdsRegistrationTimeOut.add(queryId);
+    }
 }