src/main/java/com/google/devtools/build/lib/worker/WorkerProcessMetricsCollector.java - bazel - Git at Google

 // Copyright 2022 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.worker;

 import static com.google.common.collect.ImmutableList.toImmutableList;

 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 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.util.OS;
 import java.time.Instant;
 import java.util.ArrayList;
 import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;

 /** Collects and populates system metrics about persistent workers. */
 public class WorkerProcessMetricsCollector {

   /** The metrics collector (a static singleton instance). Inactive by default. */
   private static final WorkerProcessMetricsCollector instance = new WorkerProcessMetricsCollector();

   private Clock clock;

   /**
    * Mapping of worker process ids to their process metrics. This contains all workers that have
    * been alive at any point during the build.
    */
   private final Map<Long, WorkerProcessMetrics> processIdToWorkerProcessMetrics =
       new ConcurrentHashMap<>();

   private WorkerProcessMetricsCollector() {}

   public static WorkerProcessMetricsCollector instance() {
     return instance;
   }

   public void setClock(Clock clock) {
     this.clock = clock;
   }

   /**
    * 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) {
     // TODO(b/181317827): Support Windows.
     if (processIds.isEmpty() || (os != OS.LINUX && os != OS.DARWIN)) {
       return PsInfoCollector.ResourceSnapshot.create(
           /* pidToMemoryInKb= */ ImmutableMap.of(), /* collectionTime= */ Instant.now());
     }

     return PsInfoCollector.instance().collectResourceUsage(processIds);
   }

   public ImmutableList<WorkerProcessMetrics> getLiveWorkerProcessMetrics() {
     return collectMetrics().stream()
         .filter(m -> !m.getStatus().isKilled())
         .collect(toImmutableList());
   }

   public ImmutableList<WorkerProcessMetrics> collectMetrics() {
     PsInfoCollector.ResourceSnapshot resourceSnapshot =
         collectMemoryUsageByPid(
             OS.getCurrent(), ImmutableSet.copyOf(processIdToWorkerProcessMetrics.keySet()));

     ImmutableMap<Long, Integer> pidToMemoryInKb = resourceSnapshot.getPidToMemoryInKb();
     Instant collectionTime = resourceSnapshot.getCollectionTime();

     ImmutableList.Builder<WorkerProcessMetrics> workerMetrics = new ImmutableList.Builder<>();
     for (Map.Entry<Long, WorkerProcessMetrics> entry : processIdToWorkerProcessMetrics.entrySet()) {
       WorkerProcessMetrics workerMetric = entry.getValue();
       Long pid = workerMetric.getProcessId();
       boolean isMeasurable = pidToMemoryInKb.containsKey(pid);

       if (!isMeasurable && workerMetric.getStatus().isKilled()) {
         // If it is not measurable and previously killed by Bazel, we don't do anything.
         workerMetrics.add(workerMetric);
         continue;
       } else if (!isMeasurable) {
         // If it is not measurable, not killed by Bazel but has executed actions, then we assume
         // that something has happened to the worker process that is not accounted for by Bazel
         // and set this to KILLED_UNKNOWN. If a separate thread comes along to update the status
         // with a more specific reason why it is killed, then we allow such an update.
         if (workerMetric.getActionsExecuted() > 0) {
           workerMetric.getStatus().maybeUpdateStatus(WorkerProcessStatus.Status.KILLED_UNKNOWN);
         }
         // We want to add the worker metric even if it is not measurable.
         workerMetrics.add(workerMetric);
         continue;
       }

       // If it is measurable, we want to update the collected metrics.
       workerMetric.addCollectedMetrics(
           /* memoryInKb= */ pidToMemoryInKb.getOrDefault(pid, 0),
           /* collectionTime= */ collectionTime);
       workerMetrics.add(workerMetric);
     }

     return workerMetrics.build();
   }

   public void onWorkerFinishExecution(long processId) {
     WorkerProcessMetrics wpm = processIdToWorkerProcessMetrics.get(processId);
     if (wpm == null) {
       return;
     }
     wpm.incrementActionsExecuted();
   }

   public static final int MAX_PUBLISHED_WORKER_METRICS = 50;

   /** Returns a prioritized and limited list of WorkerMetrics to be published to the BEP. */
   public ImmutableList<WorkerMetrics> getPublishedWorkerMetrics() {
     return collectMetrics().stream()
         .map(WorkerProcessMetrics::toProto)
         .sorted(new WorkerMetricsPublishComparator())
         .limit(MAX_PUBLISHED_WORKER_METRICS)
         .sorted(Comparator.comparingInt(m -> m.getWorkerIdsList().get(0)))
         .collect(toImmutableList());
   }

   public static ImmutableList<WorkerMetrics> limitWorkerMetricsToPublish(
       ImmutableList<WorkerMetrics> metrics, int limit) {
     return metrics.stream()
         .sorted(new WorkerMetricsPublishComparator())
         .limit(limit)
         .sorted(Comparator.comparingInt(m -> m.getWorkerIdsList().get(0)))
         .collect(toImmutableList());
   }

   /**
    * Because we log all worker processes that have been alive at any point during the build, the
    * size of this list might grow out of hand if there is some issue with the build (e.g.
    * kill-create cycles). As such, we enforce rules to prioritize WorkerMetrics before limiting: (1)
    * Prioritize WorkerStatuses ALIVE, then KILLED_DUE_TO_MEMORY_PRESSURE, then all remaining worker
    * statuses. (2) Then prioritize by decreasing memory usage and (3) limit to a fixed number.
    */
   @VisibleForTesting
   public static class WorkerMetricsPublishComparator implements Comparator<WorkerMetrics> {

     private int getWorkerStatusPriority(WorkerMetrics.WorkerStatus status) {
       // Lower value is prioritized.
       if (status == WorkerStatus.ALIVE) {
         return 0;
       } else if (status == WorkerStatus.KILLED_DUE_TO_MEMORY_PRESSURE) {
         return 1;
       }
       return 2;
     }

     @Override
     public int compare(WorkerMetrics m1, WorkerMetrics m2) {
       int s1 = getWorkerStatusPriority(m1.getWorkerStatus());
       int s2 = getWorkerStatusPriority(m2.getWorkerStatus());
       if (s1 != s2) {
         return Integer.compare(s1, s2);
       }
       return Integer.compare(
           m2.getWorkerStats(0).getWorkerMemoryInKb(), m1.getWorkerStats(0).getWorkerMemoryInKb());
     }
   }

   public ImmutableList<WorkerMetrics> getLiveWorkerMetrics() {
     return getLiveWorkerProcessMetrics().stream()
         .map(WorkerProcessMetrics::toProto)
         .collect(toImmutableList());
   }

   public void clear() {
     processIdToWorkerProcessMetrics.clear();
   }

   @VisibleForTesting
   Map<Long, WorkerProcessMetrics> getProcessIdToWorkerProcessMetrics() {
     return processIdToWorkerProcessMetrics;
   }

   /**
    * Initializes workerIdToWorkerProperties for workers. If worker metrics already exists for this
    * worker, only updates the last call time and maybe adds the multiplex worker id.
    */
   public synchronized void registerWorker(
       int workerId,
       long processId,
       WorkerProcessStatus status,
       String mnemonic,
       boolean isMultiplex,
       boolean isSandboxed,
       int workerKeyHash) {
     WorkerProcessMetrics workerMetric =
         processIdToWorkerProcessMetrics.computeIfAbsent(
             processId,
             (pid) ->
                 new WorkerProcessMetrics(
                     workerId,
                     processId,
                     status,
                     mnemonic,
                     isMultiplex,
                     isSandboxed,
                     workerKeyHash));

     workerMetric.setLastCallTime(Instant.ofEpochMilli(clock.currentTimeMillis()));
     workerMetric.maybeAddWorkerId(workerId, status);
   }

   /** Removes all WorkerProcessMetrics that were marked as killed. */
   public void clearKilledWorkerProcessMetrics() {
     List<Long> pidsToRemove = new ArrayList<>();
     for (Map.Entry<Long, WorkerProcessMetrics> entry : processIdToWorkerProcessMetrics.entrySet()) {
       if (entry.getValue().getStatus().isKilled()) {
         pidsToRemove.add(entry.getKey());
       }
     }
     processIdToWorkerProcessMetrics.keySet().removeAll(pidsToRemove);
   }

   /** To reset states in each WorkerProcessMetric before each command where applicable. */
   public void beforeCommand() {
     processIdToWorkerProcessMetrics.values().forEach(m -> m.onBeforeCommand());
   }

   // TODO(b/238416583) Add deregister function
 }
	// Copyright 2022 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.worker;

	import static com.google.common.collect.ImmutableList.toImmutableList;

	import com.google.common.annotations.VisibleForTesting;
	import com.google.common.collect.ImmutableList;
	import com.google.common.collect.ImmutableMap;
	import com.google.common.collect.ImmutableSet;
	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.util.OS;
	import java.time.Instant;
	import java.util.ArrayList;
	import java.util.Comparator;
	import java.util.List;
	import java.util.Map;
	import java.util.concurrent.ConcurrentHashMap;

	/** Collects and populates system metrics about persistent workers. */
	public class WorkerProcessMetricsCollector {

	/** The metrics collector (a static singleton instance). Inactive by default. */
	private static final WorkerProcessMetricsCollector instance = new WorkerProcessMetricsCollector();

	private Clock clock;

	/**
	* Mapping of worker process ids to their process metrics. This contains all workers that have
	* been alive at any point during the build.
	*/
	private final Map<Long, WorkerProcessMetrics> processIdToWorkerProcessMetrics =
	new ConcurrentHashMap<>();

	private WorkerProcessMetricsCollector() {}

	public static WorkerProcessMetricsCollector instance() {
	return instance;
	}

	public void setClock(Clock clock) {
	this.clock = clock;
	}

	/**
	* 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) {
	// TODO(b/181317827): Support Windows.
	if (processIds.isEmpty() \|\| (os != OS.LINUX && os != OS.DARWIN)) {
	return PsInfoCollector.ResourceSnapshot.create(
	/* pidToMemoryInKb= / ImmutableMap.of(), / collectionTime= */ Instant.now());
	}

	return PsInfoCollector.instance().collectResourceUsage(processIds);
	}

	public ImmutableList<WorkerProcessMetrics> getLiveWorkerProcessMetrics() {
	return collectMetrics().stream()
	.filter(m -> !m.getStatus().isKilled())
	.collect(toImmutableList());
	}

	public ImmutableList<WorkerProcessMetrics> collectMetrics() {
	PsInfoCollector.ResourceSnapshot resourceSnapshot =
	collectMemoryUsageByPid(
	OS.getCurrent(), ImmutableSet.copyOf(processIdToWorkerProcessMetrics.keySet()));

	ImmutableMap<Long, Integer> pidToMemoryInKb = resourceSnapshot.getPidToMemoryInKb();
	Instant collectionTime = resourceSnapshot.getCollectionTime();

	ImmutableList.Builder<WorkerProcessMetrics> workerMetrics = new ImmutableList.Builder<>();
	for (Map.Entry<Long, WorkerProcessMetrics> entry : processIdToWorkerProcessMetrics.entrySet()) {
	WorkerProcessMetrics workerMetric = entry.getValue();
	Long pid = workerMetric.getProcessId();
	boolean isMeasurable = pidToMemoryInKb.containsKey(pid);

	if (!isMeasurable && workerMetric.getStatus().isKilled()) {
	// If it is not measurable and previously killed by Bazel, we don't do anything.
	workerMetrics.add(workerMetric);
	continue;
	} else if (!isMeasurable) {
	// If it is not measurable, not killed by Bazel but has executed actions, then we assume
	// that something has happened to the worker process that is not accounted for by Bazel
	// and set this to KILLED_UNKNOWN. If a separate thread comes along to update the status
	// with a more specific reason why it is killed, then we allow such an update.
	if (workerMetric.getActionsExecuted() > 0) {
	workerMetric.getStatus().maybeUpdateStatus(WorkerProcessStatus.Status.KILLED_UNKNOWN);
	}
	// We want to add the worker metric even if it is not measurable.
	workerMetrics.add(workerMetric);
	continue;
	}

	// If it is measurable, we want to update the collected metrics.
	workerMetric.addCollectedMetrics(
	/* memoryInKb= */ pidToMemoryInKb.getOrDefault(pid, 0),
	/* collectionTime= */ collectionTime);
	workerMetrics.add(workerMetric);
	}

	return workerMetrics.build();
	}

	public void onWorkerFinishExecution(long processId) {
	WorkerProcessMetrics wpm = processIdToWorkerProcessMetrics.get(processId);
	if (wpm == null) {
	return;
	}
	wpm.incrementActionsExecuted();
	}

	public static final int MAX_PUBLISHED_WORKER_METRICS = 50;

	/** Returns a prioritized and limited list of WorkerMetrics to be published to the BEP. */
	public ImmutableList<WorkerMetrics> getPublishedWorkerMetrics() {
	return collectMetrics().stream()
	.map(WorkerProcessMetrics::toProto)
	.sorted(new WorkerMetricsPublishComparator())
	.limit(MAX_PUBLISHED_WORKER_METRICS)
	.sorted(Comparator.comparingInt(m -> m.getWorkerIdsList().get(0)))
	.collect(toImmutableList());
	}

	public static ImmutableList<WorkerMetrics> limitWorkerMetricsToPublish(
	ImmutableList<WorkerMetrics> metrics, int limit) {
	return metrics.stream()
	.sorted(new WorkerMetricsPublishComparator())
	.limit(limit)
	.sorted(Comparator.comparingInt(m -> m.getWorkerIdsList().get(0)))
	.collect(toImmutableList());
	}

	/**
	* Because we log all worker processes that have been alive at any point during the build, the
	* size of this list might grow out of hand if there is some issue with the build (e.g.
	* kill-create cycles). As such, we enforce rules to prioritize WorkerMetrics before limiting: (1)
	* Prioritize WorkerStatuses ALIVE, then KILLED_DUE_TO_MEMORY_PRESSURE, then all remaining worker
	* statuses. (2) Then prioritize by decreasing memory usage and (3) limit to a fixed number.
	*/
	@VisibleForTesting
	public static class WorkerMetricsPublishComparator implements Comparator<WorkerMetrics> {

	private int getWorkerStatusPriority(WorkerMetrics.WorkerStatus status) {
	// Lower value is prioritized.
	if (status == WorkerStatus.ALIVE) {
	return 0;
	} else if (status == WorkerStatus.KILLED_DUE_TO_MEMORY_PRESSURE) {
	return 1;
	}
	return 2;
	}

	@Override
	public int compare(WorkerMetrics m1, WorkerMetrics m2) {
	int s1 = getWorkerStatusPriority(m1.getWorkerStatus());
	int s2 = getWorkerStatusPriority(m2.getWorkerStatus());
	if (s1 != s2) {
	return Integer.compare(s1, s2);
	}
	return Integer.compare(
	m2.getWorkerStats(0).getWorkerMemoryInKb(), m1.getWorkerStats(0).getWorkerMemoryInKb());
	}
	}

	public ImmutableList<WorkerMetrics> getLiveWorkerMetrics() {
	return getLiveWorkerProcessMetrics().stream()
	.map(WorkerProcessMetrics::toProto)
	.collect(toImmutableList());
	}

	public void clear() {
	processIdToWorkerProcessMetrics.clear();
	}

	@VisibleForTesting
	Map<Long, WorkerProcessMetrics> getProcessIdToWorkerProcessMetrics() {
	return processIdToWorkerProcessMetrics;
	}

	/**
	* Initializes workerIdToWorkerProperties for workers. If worker metrics already exists for this
	* worker, only updates the last call time and maybe adds the multiplex worker id.
	*/
	public synchronized void registerWorker(
	int workerId,
	long processId,
	WorkerProcessStatus status,
	String mnemonic,
	boolean isMultiplex,
	boolean isSandboxed,
	int workerKeyHash) {
	WorkerProcessMetrics workerMetric =
	processIdToWorkerProcessMetrics.computeIfAbsent(
	processId,
	(pid) ->
	new WorkerProcessMetrics(
	workerId,
	processId,
	status,
	mnemonic,
	isMultiplex,
	isSandboxed,
	workerKeyHash));

	workerMetric.setLastCallTime(Instant.ofEpochMilli(clock.currentTimeMillis()));
	workerMetric.maybeAddWorkerId(workerId, status);
	}

	/** Removes all WorkerProcessMetrics that were marked as killed. */
	public void clearKilledWorkerProcessMetrics() {
	List<Long> pidsToRemove = new ArrayList<>();
	for (Map.Entry<Long, WorkerProcessMetrics> entry : processIdToWorkerProcessMetrics.entrySet()) {
	if (entry.getValue().getStatus().isKilled()) {
	pidsToRemove.add(entry.getKey());
	}
	}
	processIdToWorkerProcessMetrics.keySet().removeAll(pidsToRemove);
	}

	/** To reset states in each WorkerProcessMetric before each command where applicable. */
	public void beforeCommand() {
	processIdToWorkerProcessMetrics.values().forEach(m -> m.onBeforeCommand());
	}

	// TODO(b/238416583) Add deregister function
	}