Make PlanOptimizersStatsCollector thread-safe

The class can be accessed concurrently when query is cancelled during
planning. The planning thread is interrupted but may continue running,
and query completion even contains information got from
`PlanOptimizersStatsCollector`.

core/trino-main/src/main/java/io/trino/execution/querystats/PlanOptimizersStatsCollector.java

@@ -17,16 +17,19 @@
 import io.trino.sql.planner.iterative.Rule;
 import io.trino.sql.planner.optimizations.PlanOptimizer;
 
+import javax.annotation.concurrent.ThreadSafe;
+
 import java.util.Comparator;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
 
+@ThreadSafe
 public class PlanOptimizersStatsCollector
 {
-    private final Map<Class<?>, QueryPlanOptimizerStats> stats = new HashMap<>();
+    private final Map<Class<?>, QueryPlanOptimizerStats> stats = new ConcurrentHashMap<>();
     private final int queryReportedRuleStatsLimit;
 
     public PlanOptimizersStatsCollector(int queryReportedRuleStatsLimit)

core/trino-main/src/main/java/io/trino/execution/querystats/QueryPlanOptimizerStats.java

@@ -15,13 +15,18 @@
 
 import io.trino.spi.eventlistener.QueryPlanOptimizerStatistics;
 
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.util.concurrent.atomic.AtomicLong;
+
+@ThreadSafe
 public class QueryPlanOptimizerStats
 {
     private final String rule;
-    private long invocations;
-    private long applied;
-    private long totalTime;
-    private long failures;
+    private final AtomicLong invocations = new AtomicLong();
+    private final AtomicLong applied = new AtomicLong();
+    private final AtomicLong totalTime = new AtomicLong();
+    private final AtomicLong failures = new AtomicLong();
 
     public QueryPlanOptimizerStats(String rule)
     {
@@ -31,16 +36,16 @@ public QueryPlanOptimizerStats(String rule)
     public void record(long nanos, boolean applied)
     {
         if (applied) {
-            this.applied += 1;
+            this.applied.incrementAndGet();
         }
 
-        invocations += 1;
-        totalTime += nanos;
+        invocations.incrementAndGet();
+        totalTime.addAndGet(nanos);
     }
 
     public void recordFailure()
     {
-        failures += 1;
+        failures.incrementAndGet();
     }
 
     public String getRule()
@@ -50,35 +55,35 @@ public String getRule()
 
     public long getInvocations()
     {
-        return invocations;
+        return invocations.get();
     }
 
     public long getApplied()
     {
-        return applied;
+        return applied.get();
     }
 
     public long getFailures()
     {
-        return failures;
+        return failures.get();
     }
 
     public long getTotalTime()
     {
-        return totalTime;
+        return totalTime.get();
     }
 
     public QueryPlanOptimizerStatistics snapshot(String rule)
     {
-        return new QueryPlanOptimizerStatistics(rule, invocations, applied, totalTime, failures);
+        return new QueryPlanOptimizerStatistics(rule, invocations.get(), applied.get(), totalTime.get(), failures.get());
     }
 
     public QueryPlanOptimizerStats merge(QueryPlanOptimizerStats other)
     {
-        invocations += other.getInvocations();
-        applied += other.getApplied();
-        failures += other.getFailures();
-        totalTime += other.getTotalTime();
+        invocations.addAndGet(other.getInvocations());
+        applied.addAndGet(other.getApplied());
+        failures.addAndGet(other.getFailures());
+        totalTime.addAndGet(other.getTotalTime());
 
         return this;
     }