Move ResourceSnapshot out of PsInfoCollector (we'll be using creating this snapshot with cgroups which has nothing to do with ps) and also only collect resources from alive workers.
PiperOrigin-RevId: 601125773
Change-Id: I3dc2c2d62fd15caf3d58b13ac1a96ddd6471aee6
diff --git a/src/main/java/com/google/devtools/build/lib/metrics/BUILD b/src/main/java/com/google/devtools/build/lib/metrics/BUILD
index cd0efda..02c0450 100644
--- a/src/main/java/com/google/devtools/build/lib/metrics/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/metrics/BUILD
@@ -88,7 +88,10 @@
java_library(
name = "ps_info_collector",
- srcs = ["PsInfoCollector.java"],
+ srcs = [
+ "PsInfoCollector.java",
+ "ResourceSnapshot.java",
+ ],
deps = [
"//src/main/java/com/google/devtools/build/lib/clock",
"//third_party:auto_value",
diff --git a/src/main/java/com/google/devtools/build/lib/metrics/PsInfoCollector.java b/src/main/java/com/google/devtools/build/lib/metrics/PsInfoCollector.java
index 952c256..00cd48f 100644
--- a/src/main/java/com/google/devtools/build/lib/metrics/PsInfoCollector.java
+++ b/src/main/java/com/google/devtools/build/lib/metrics/PsInfoCollector.java
@@ -190,19 +190,4 @@
pidToPsInfo, pidToChildrenPsInfo, collectionTime);
}
}
-
- /** Contains resource consumption info per process tree. */
- @AutoValue
- public abstract static class ResourceSnapshot {
- /** Overall memory consumption by all descendant processes including initial process. */
- public abstract ImmutableMap<Long, Integer> getPidToMemoryInKb();
-
- /** Time when this snapshot was collected. */
- public abstract Instant getCollectionTime();
-
- public static ResourceSnapshot create(
- ImmutableMap<Long, Integer> pidToMemoryInKb, Instant collectionTime) {
- return new AutoValue_PsInfoCollector_ResourceSnapshot(pidToMemoryInKb, collectionTime);
- }
- }
}
diff --git a/src/main/java/com/google/devtools/build/lib/metrics/ResourceSnapshot.java b/src/main/java/com/google/devtools/build/lib/metrics/ResourceSnapshot.java
new file mode 100644
index 0000000..99009f7
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/metrics/ResourceSnapshot.java
@@ -0,0 +1,38 @@
+// Copyright 2024 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.devtools.build.lib.metrics;
+
+import com.google.auto.value.AutoValue;
+import com.google.common.collect.ImmutableMap;
+import java.time.Instant;
+
+/** Contains a snapshot of the resource usage of multiple processes. */
+@AutoValue
+public abstract class ResourceSnapshot {
+ /** Overall memory consumption by all descendant processes including initial process. */
+ public abstract ImmutableMap<Long, Integer> getPidToMemoryInKb();
+
+ /** Time when this snapshot was collected. */
+ public abstract Instant getCollectionTime();
+
+ public static ResourceSnapshot create(
+ ImmutableMap<Long, Integer> pidToMemoryInKb, Instant collectionTime) {
+ return new AutoValue_ResourceSnapshot(pidToMemoryInKb, collectionTime);
+ }
+
+ public static ResourceSnapshot createEmpty() {
+ return new AutoValue_ResourceSnapshot(ImmutableMap.of(), Instant.now());
+ }
+}
diff --git a/src/main/java/com/google/devtools/build/lib/worker/WorkerProcessMetricsCollector.java b/src/main/java/com/google/devtools/build/lib/worker/WorkerProcessMetricsCollector.java
index 3636cff..ba998a2 100644
--- a/src/main/java/com/google/devtools/build/lib/worker/WorkerProcessMetricsCollector.java
+++ b/src/main/java/com/google/devtools/build/lib/worker/WorkerProcessMetricsCollector.java
@@ -15,6 +15,7 @@
package com.google.devtools.build.lib.worker;
import static com.google.common.collect.ImmutableList.toImmutableList;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ImmutableList;
@@ -24,6 +25,7 @@
import com.google.devtools.build.lib.buildeventstream.BuildEventStreamProtos.BuildMetrics.WorkerMetrics.WorkerStatus;
import com.google.devtools.build.lib.clock.Clock;
import com.google.devtools.build.lib.metrics.PsInfoCollector;
+import com.google.devtools.build.lib.metrics.ResourceSnapshot;
import com.google.devtools.build.lib.util.OS;
import java.time.Instant;
import java.util.ArrayList;
@@ -57,18 +59,29 @@
this.clock = clock;
}
+ ResourceSnapshot collectResourceUsage() {
+ // Only collect for process we know are alive.
+ ImmutableSet<Long> alivePids =
+ processIdToWorkerProcessMetrics.entrySet().stream()
+ .filter(e -> !e.getValue().getStatus().isKilled())
+ .map(e -> e.getKey())
+ .collect(toImmutableSet());
+ return collectResourceUsage(OS.getCurrent(), alivePids);
+ }
+
/**
* Collects memory usage of all ancestors of processes by pid. If a pid does not allow collecting
* memory usage, it is silently ignored.
*/
- PsInfoCollector.ResourceSnapshot collectMemoryUsageByPid(OS os, ImmutableSet<Long> processIds) {
+ ResourceSnapshot collectResourceUsage(OS os, ImmutableSet<Long> processIds) {
// TODO(b/181317827): Support Windows.
- if (processIds.isEmpty() || (os != OS.LINUX && os != OS.DARWIN)) {
- return PsInfoCollector.ResourceSnapshot.create(
- /* pidToMemoryInKb= */ ImmutableMap.of(), /* collectionTime= */ Instant.now());
+ if (processIds.isEmpty()) {
+ return ResourceSnapshot.createEmpty();
}
-
- return PsInfoCollector.instance().collectResourceUsage(processIds);
+ if (os == OS.LINUX || os == OS.DARWIN) {
+ return PsInfoCollector.instance().collectResourceUsage(processIds);
+ }
+ return ResourceSnapshot.createEmpty();
}
public ImmutableList<WorkerProcessMetrics> getLiveWorkerProcessMetrics() {
@@ -78,9 +91,7 @@
}
public ImmutableList<WorkerProcessMetrics> collectMetrics() {
- PsInfoCollector.ResourceSnapshot resourceSnapshot =
- collectMemoryUsageByPid(
- OS.getCurrent(), ImmutableSet.copyOf(processIdToWorkerProcessMetrics.keySet()));
+ ResourceSnapshot resourceSnapshot = collectResourceUsage();
ImmutableMap<Long, Integer> pidToMemoryInKb = resourceSnapshot.getPidToMemoryInKb();
Instant collectionTime = resourceSnapshot.getCollectionTime();
