Allow register CounterSeriesCollector

PiperOrigin-RevId: 675920887 Change-Id: Ie57ff05818b8c58319513c66573143c034819927
bazelbuild · Sep 18, 2024 · ae46abf · ae46abf
1 parent 538136e
commit ae46abf
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 5 deletions.
diff --git a/src/main/java/com/google/devtools/build/lib/profiler/CollectLocalResourceUsage.java b/src/main/java/com/google/devtools/build/lib/profiler/CollectLocalResourceUsage.java
@@ -48,6 +48,7 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.TimeUnit;
 import java.util.function.BiConsumer;
 import javax.annotation.Nullable;
@@ -69,9 +70,8 @@ public class CollectLocalResourceUsage implements LocalResourceCollector {
   private volatile boolean stopLocalUsageCollection;
   private volatile boolean profilingStarted;
 
-  @GuardedBy("this")
-  @Nullable
-  private List<CounterSeriesCollector> collectors;
+  private final ConcurrentLinkedQueue<CounterSeriesCollector> collectors =
+      new ConcurrentLinkedQueue<>();
 
   @GuardedBy("this")
   @Nullable
@@ -122,6 +122,11 @@ public void start() {
     collector.start();
   }
 
+  @Override
+  public void registerCounterSeriesCollector(CounterSeriesCollector collector) {
+    collectors.add(collector);
+  }
+
   private List<CounterSeriesCollector> createLocalCollectors() {
     OperatingSystemMXBean osBean =
         (OperatingSystemMXBean) ManagementFactory.getOperatingSystemMXBean();
@@ -165,8 +170,8 @@ private class Collector extends Thread {
 
     @Override
     public void run() {
+      collectors.addAll(createLocalCollectors());
       synchronized (CollectLocalResourceUsage.this) {
-        collectors = new ArrayList<>(createLocalCollectors());
         timeSeries = new LinkedHashMap<>();
       }
 
@@ -236,7 +241,7 @@ synchronized void logCollectedData() {
       profiler.logCounters(stackedCounters.buildOrThrow(), profileStart, BUCKET_DURATION);
     }
 
-    collectors = null;
+    collectors.clear();
     timeSeries = null;
   }
 

diff --git a/src/main/java/com/google/devtools/build/lib/profiler/LocalResourceCollector.java b/src/main/java/com/google/devtools/build/lib/profiler/LocalResourceCollector.java
@@ -14,12 +14,17 @@
 
 package com.google.devtools.build.lib.profiler;
 
+import com.google.devtools.build.lib.profiler.Profiler.CounterSeriesCollector;
+
 /** Collector used by the profiler to collect local resources. */
 public interface LocalResourceCollector {
 
   /** Start collecting. */
   void start();
 
+  /** Register a counter series collector. */
+  void registerCounterSeriesCollector(CounterSeriesCollector collector);
+
   /** Stop collecting. */
   void stop();
 }
diff --git a/src/main/java/com/google/devtools/build/lib/profiler/Profiler.java b/src/main/java/com/google/devtools/build/lib/profiler/Profiler.java
@@ -640,6 +640,14 @@ public interface CounterSeriesCollector {
     void collect(double deltaNanos, BiConsumer<CounterSeriesTask, Double> consumer);
   }
 
+  public void registerCounterSeriesCollector(CounterSeriesCollector collector) {
+    if (!isActive()) {
+      return;
+    }
+
+    localResourceCollector.registerCounterSeriesCollector(collector);
+  }
+
   /** Adds a whole action count series to the writer bypassing histogram and subtask creation. */
   public void logCounters(
       Map<CounterSeriesTask, double[]> counterSeriesMap,