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();
diff --git a/src/test/java/com/google/devtools/build/lib/metrics/PsInfoCollectorTest.java b/src/test/java/com/google/devtools/build/lib/metrics/PsInfoCollectorTest.java
index d12783a..3372fdc 100644
--- a/src/test/java/com/google/devtools/build/lib/metrics/PsInfoCollectorTest.java
+++ b/src/test/java/com/google/devtools/build/lib/metrics/PsInfoCollectorTest.java
@@ -87,7 +87,7 @@
ImmutableSet<Long> pids = ImmutableSet.of(1L, 2L, 5L, 6L);
when(spyCollector.collectDataFromPs()).thenReturn(psInfos);
- PsInfoCollector.ResourceSnapshot resourceSnapshot = spyCollector.collectResourceUsage(pids);
+ ResourceSnapshot resourceSnapshot = spyCollector.collectResourceUsage(pids);
ImmutableMap<Long, Integer> expectedMemoryUsageByPid =
ImmutableMap.of(1L, 3216 + 1234 + 2345 + 3456, 2L, 4232 + 1001 + 1032, 5L, 40000);
diff --git a/src/test/java/com/google/devtools/build/lib/worker/WorkerProcessMetricsCollectorTest.java b/src/test/java/com/google/devtools/build/lib/worker/WorkerProcessMetricsCollectorTest.java
index 743c555..bdfe29c 100644
--- a/src/test/java/com/google/devtools/build/lib/worker/WorkerProcessMetricsCollectorTest.java
+++ b/src/test/java/com/google/devtools/build/lib/worker/WorkerProcessMetricsCollectorTest.java
@@ -25,7 +25,7 @@
import com.google.devtools.build.lib.buildeventstream.BuildEventStreamProtos.BuildMetrics.WorkerMetrics;
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.worker.WorkerProcessMetricsCollector.WorkerMetricsPublishComparator;
import com.google.devtools.build.lib.worker.WorkerProcessStatus.Status;
import java.time.Instant;
@@ -310,9 +310,8 @@
PROCESS_ID_1, 1234,
PROCESS_ID_2, 2345);
Instant collectionTime = DEFAULT_CLOCK_START_INSTANT.plusSeconds(10);
- PsInfoCollector.ResourceSnapshot resourceSnapshot =
- PsInfoCollector.ResourceSnapshot.create(memoryUsageMap, collectionTime);
- doReturn(resourceSnapshot).when(spyCollector).collectMemoryUsageByPid(any(), any());
+ ResourceSnapshot resourceSnapshot = ResourceSnapshot.create(memoryUsageMap, collectionTime);
+ doReturn(resourceSnapshot).when(spyCollector).collectResourceUsage(any(), any());
clock.setTime(collectionTime.toEpochMilli());
ImmutableList<WorkerProcessMetrics> metrics = spyCollector.collectMetrics();
