src/main/java/com/google/devtools/build/skyframe/SkyframeFocuser.java - bazel - Git at Google

 // Copyright 2023 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.skyframe;

 import static java.util.concurrent.TimeUnit.MINUTES;

 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Sets;
 import com.google.devtools.build.lib.concurrent.AbstractQueueVisitor;
 import com.google.devtools.build.lib.concurrent.ErrorClassifier;
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.events.EventHandler;
 import com.google.devtools.build.lib.profiler.Profiler;
 import com.google.devtools.build.lib.profiler.SilentCloseable;
 import java.util.Collection;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.Function;

 /**
  * SkyframeFocuser is a minimizing optimizer (i.e. garbage collector) for the Skyframe graph, based
  * on a set of known inputs known as a working set, while ensuring correct incremental builds.
  *
  * <p>This is also a subclass of {@link AbstractQueueVisitor} to take advantage of highly
  * parallelizable operations over the Skyframe graph.
  */
 public final class SkyframeFocuser extends AbstractQueueVisitor {

   // The in-memory Skyframe graph
   private final InMemoryGraph graph;

   // Event handler to report stats during focusing
   private final EventHandler eventHandler;

   private SkyframeFocuser(InMemoryGraph graph, EventHandler eventHandler) {
     super(
         /* parallelism= */ Runtime.getRuntime().availableProcessors(),
         /* keepAliveTime= */ 2,
         MINUTES,
         ExceptionHandlingMode.FAIL_FAST,
         /* poolName= */ "skyframe-focuser",
         ErrorClassifier.DEFAULT);
     this.graph = graph;
     this.eventHandler = eventHandler;
   }

   /**
    * Minimize the Skyframe graph by traverse it to prune nodes and edges that are not necessary for
    * the build correctness of a working set of files. The graph focusing algorithm pseudocode is as
    * follows.
    *
    * <ol>
    *   <li>Mark all the leafs and their transitive rdeps. For each marked node, also mark all their
    *       direct dependencies. An injectable function can also mark additional nodes reachable from
    *       the node itself.
    *   <li>For each marked node, remove all direct deps edges. Also remove all rdep edges unless
    *       they point to a rdep that should be kept. This creates the "flattened verification set".
    * </ol>
    *
    * @param graph the in-memory graph to operate on
    * @param eventHandler handler to report events during focusing
    * @param roots the SkyKeys of the roots to be kept, i.e. the top level keys.
    * @param leafs the SkyKeys of the leafs to be kept. This is the "working set".
    * @param additionalDepsToKeep by default, all direct deps of rdeps are kept. this function is
    *     applied on all direct deps, in case there are optimizations elsewhere in the skyframe
    *     implementation that reads transitive nodes without specifying a dependency on them.
    * @return the set of kept SkyKeys in the in-memory graph, categorized by deps and rdeps.
    */
   public static FocusResult focus(
       InMemoryGraph graph,
       EventHandler eventHandler,
       Set<SkyKey> roots,
       Set<SkyKey> leafs,
       Function<SkyKey, Set<SkyKey>> additionalDepsToKeep)
       throws InterruptedException {
     SkyframeFocuser focuser = new SkyframeFocuser(graph, eventHandler);
     return focuser.run(roots, leafs, additionalDepsToKeep);
   }

   /**
    * The set of SkyKeys kept after focusing. The actual change is done in place with the in-memory
    * graph.
    */
   public static class FocusResult {

     private final ImmutableSet<SkyKey> rdeps;
     private final ImmutableSet<SkyKey> deps;

     private FocusResult(ImmutableSet<SkyKey> rdeps, ImmutableSet<SkyKey> deps) {
       this.rdeps = rdeps;
       this.deps = deps;
     }

     /**
      * Returns the set of SkyKeys that are in the dependencies of all roots, and rdeps from the
      * leafs. May contain transitive dependencies, in cases where certain functions use them without
      * establishing a Skyframe dependency.
      */
     public ImmutableSet<SkyKey> getDeps() {
       return deps;
     }

     /** Returns the set of SkyKeys that are in the reverse dependencies of the leafs. */
     public ImmutableSet<SkyKey> getRdeps() {
       return rdeps;
     }
   }

   /**
    * NodeVisitor is parallelizable graph visitor that's applied transitively from leafs to the
    * roots, while marking rdeps and all direct deps of those rdeps to be kept by SkyframeFocuser.
    */
   private class NodeVisitor implements Runnable {

     // The SkyKey that this NodeVisitor is responsible for.
     private final SkyKey key;

     // Threadsafe set of keys that depend on this key. May be modified by multiple NodeVisitors
     // concurrently.
     private final Set<SkyKey> keptRdeps;

     // Threadsafe set of keys that this key depends on. May be modified by multiple NodeVisitors
     // concurrently.
     private final Set<SkyKey> keptDeps;

     // See javadoc for the additionalDepsToKeep parameter of {@code SkyframeFocuser#focus}.
     private final Function<SkyKey, Set<SkyKey>> additionalDepsToKeep;

     protected NodeVisitor(
         SkyKey key,
         Set<SkyKey> keptRdeps,
         Set<SkyKey> keptDeps,
         Function<SkyKey, Set<SkyKey>> additionalDepsToKeep) {
       this.key = key;
       this.keptRdeps = keptRdeps;
       this.keptDeps = keptDeps;
       this.additionalDepsToKeep = additionalDepsToKeep;
     }

     @Override
     public void run() {
       InMemoryNodeEntry nodeEntry = graph.getIfPresent(key);
       if (nodeEntry == null) {
         // Throwing may be too strong if the working set is loosely defined, e.g. an entire
         // directory instead of a single file, and there are some files in the directory that
         // are not in the TC of the roots.
         //
         // TODO: b/312819241 - handle this gracefully without throwing.
         throw new IllegalStateException("nodeEntry not found for: " + key.getCanonicalName());
       }

       if (!nodeEntry.isDone()) {
         // TODO: b/312819241 - handle this gracefully without throwing.
         throw new IllegalStateException("nodeEntry not done: " + key.getCanonicalName());
       }

       for (SkyKey rdep : nodeEntry.getReverseDepsForDoneEntry()) {
         if (!keptRdeps.add(rdep)) {
           // Memoization. Already processed.
           continue;
         }

         // Queue a traversal up the graph. This will not create duplicate NodeVisitors on the
         // same rdep due to the atomic keptRdeps.add check above.
         execute(new NodeVisitor(rdep, keptRdeps, keptDeps, additionalDepsToKeep));
       }

       for (SkyKey dep : nodeEntry.getDirectDeps()) {
         if (!keptDeps.add(dep)) {
           // Memoization. Already processed.
           continue;
         }

         // This is necessary to keep the action inputs encapsulated by a NestedSet. Otherwise,
         // those inputs will be missing.
         //
         // TODO: b/312819241 - move SkyframeFocuser from build.skyframe to build.lib.skyframe so
         // that we can do the check directly without using an injected Function.
         execute(() -> keptDeps.addAll(additionalDepsToKeep.apply(dep)));
       }
     }
   }

   /** Entry point of the Skyframe garbage collection algorithm. */
   private FocusResult run(
       Set<SkyKey> roots, Set<SkyKey> leafs, Function<SkyKey, Set<SkyKey>> additionalDepsToKeep)
       throws InterruptedException {

     Set<SkyKey> keptDeps = Sets.newConcurrentHashSet();
     Set<SkyKey> keptRdeps = Sets.newConcurrentHashSet();

     // All leafs are automatically considered as rdeps.
     keptRdeps.addAll(leafs);

     // All roots are automatically considered as deps.
     //
     // Some roots are re-evaluated on every build. These roots may not be in the reverse TC
     // of leafs (working set), but may influence how the working set is evaluated (e.g. platform
     // mapping). If we remove them from the graph, those keys may be re-evaluated anyway (along with
     // their TC) on subsequent invocations, leading to wasted compute and RAM.
     // The exercise of ensuring which roots should be kept is left to the caller of
     // this function, but we ensure that all specified ones are kept here.
     keptDeps.addAll(roots);

     try (SilentCloseable c = Profiler.instance().profile("focus.mark")) {
       // Start traversal from leafs.
       for (SkyKey leaf : leafs) {
         execute(new NodeVisitor(leaf, keptRdeps, keptDeps, additionalDepsToKeep));
       }
       awaitQuiescenceWithoutShutdown(true);
     }

     // Keep the rdeps transitive closure from leafs distinct from the deps.
     keptDeps.removeAll(keptRdeps);

     eventHandler.handle(
         Event.info("Nodes in reverse transitive closure from leafs: " + keptRdeps.size()));
     eventHandler.handle(
         Event.info("Nodes in direct deps of reverse transitive closure: " + keptDeps.size()));

     try (SilentCloseable c = Profiler.instance().profile("focus.sweep_nodes")) {
       Set<SkyKey> toKeep = Sets.union(keptDeps, keptRdeps);
       graph.parallelForEach(
           inMemoryNodeEntry -> {
             SkyKey key = inMemoryNodeEntry.getKey();
             if (!toKeep.contains(key)) {
               graph.remove(key);
             }
           });

       graph.shrinkNodeMap();
     }

     AtomicLong rdepEdgesBefore = new AtomicLong();
     AtomicLong rdepEdgesAfter = new AtomicLong();

     try (SilentCloseable c = Profiler.instance().profile("focus.sweep_edges")) {
       for (SkyKey key : keptDeps) {
         execute(
             () -> {
               // TODO: b/312819241 - Consider transforming IncrementalInMemoryNodeEntry only used
               // for their immutable states to an ImmutableDoneNodeEntry or
               // NonIncrementalInMemoryNodeEntry
               // for further memory savings.
               IncrementalInMemoryNodeEntry nodeEntry =
                   (IncrementalInMemoryNodeEntry) graph.getIfPresent(key);

               // No need to keep the direct deps edges of existing deps. For example:
               //
               //    B
               //  / |\
               // A C  \
               //    \ |
               //     D
               //
               // B is the root, and A is the only leaf. We can throw out the CD edge, even
               // though both C and D are still used by B. This is because no changes are expected to
               // C and D, so it's unnecessary to maintain the edges.
               nodeEntry.directDeps = GroupedDeps.EMPTY_COMPRESSED;

               // No need to keep the rdep edges of the deps if they do not point to an rdep
               // reachable (hence, dirty-able) by the working set.
               //
               // This accounts for nearly 5% of 9+GB retained heap on a large server build.
               Collection<SkyKey> existingRdeps = nodeEntry.getReverseDepsForDoneEntry();
               rdepEdgesBefore.getAndAdd(existingRdeps.size());
               int rdepEdgesKept = 0;
               for (SkyKey rdep : existingRdeps) {
                 if (keptRdeps.contains(rdep)) {
                   rdepEdgesKept++;
                 } else {
                   nodeEntry.removeReverseDep(rdep);
                 }
               }
               rdepEdgesAfter.getAndAdd(rdepEdgesKept);

               ReverseDepsUtility.consolidateData(nodeEntry);
             });
       }

       awaitQuiescence(true); // and shut down the ExecutorService.
     }

     eventHandler.handle(
         Event.info(
             String.format("Rdep edges: %s -> %s", rdepEdgesBefore.get(), rdepEdgesAfter.get())));

     return new FocusResult(ImmutableSet.copyOf(keptRdeps), ImmutableSet.copyOf(keptDeps));
   }
 }
	// Copyright 2023 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.skyframe;

	import static java.util.concurrent.TimeUnit.MINUTES;

	import com.google.common.collect.ImmutableSet;
	import com.google.common.collect.Sets;
	import com.google.devtools.build.lib.concurrent.AbstractQueueVisitor;
	import com.google.devtools.build.lib.concurrent.ErrorClassifier;
	import com.google.devtools.build.lib.events.Event;
	import com.google.devtools.build.lib.events.EventHandler;
	import com.google.devtools.build.lib.profiler.Profiler;
	import com.google.devtools.build.lib.profiler.SilentCloseable;
	import java.util.Collection;
	import java.util.Set;
	import java.util.concurrent.atomic.AtomicLong;
	import java.util.function.Function;

	/**
	* SkyframeFocuser is a minimizing optimizer (i.e. garbage collector) for the Skyframe graph, based
	* on a set of known inputs known as a working set, while ensuring correct incremental builds.
	*
	* <p>This is also a subclass of {@link AbstractQueueVisitor} to take advantage of highly
	* parallelizable operations over the Skyframe graph.
	*/
	public final class SkyframeFocuser extends AbstractQueueVisitor {

	// The in-memory Skyframe graph
	private final InMemoryGraph graph;

	// Event handler to report stats during focusing
	private final EventHandler eventHandler;

	private SkyframeFocuser(InMemoryGraph graph, EventHandler eventHandler) {
	super(
	/* parallelism= */ Runtime.getRuntime().availableProcessors(),
	/* keepAliveTime= */ 2,
	MINUTES,
	ExceptionHandlingMode.FAIL_FAST,
	/* poolName= */ "skyframe-focuser",
	ErrorClassifier.DEFAULT);
	this.graph = graph;
	this.eventHandler = eventHandler;
	}

	/**
	* Minimize the Skyframe graph by traverse it to prune nodes and edges that are not necessary for
	* the build correctness of a working set of files. The graph focusing algorithm pseudocode is as
	* follows.
	*
	* <ol>
	* <li>Mark all the leafs and their transitive rdeps. For each marked node, also mark all their
	* direct dependencies. An injectable function can also mark additional nodes reachable from
	* the node itself.
	* <li>For each marked node, remove all direct deps edges. Also remove all rdep edges unless
	* they point to a rdep that should be kept. This creates the "flattened verification set".
	* </ol>
	*
	* @param graph the in-memory graph to operate on
	* @param eventHandler handler to report events during focusing
	* @param roots the SkyKeys of the roots to be kept, i.e. the top level keys.
	* @param leafs the SkyKeys of the leafs to be kept. This is the "working set".
	* @param additionalDepsToKeep by default, all direct deps of rdeps are kept. this function is
	* applied on all direct deps, in case there are optimizations elsewhere in the skyframe
	* implementation that reads transitive nodes without specifying a dependency on them.
	* @return the set of kept SkyKeys in the in-memory graph, categorized by deps and rdeps.
	*/
	public static FocusResult focus(
	InMemoryGraph graph,
	EventHandler eventHandler,
	Set<SkyKey> roots,
	Set<SkyKey> leafs,
	Function<SkyKey, Set<SkyKey>> additionalDepsToKeep)
	throws InterruptedException {
	SkyframeFocuser focuser = new SkyframeFocuser(graph, eventHandler);
	return focuser.run(roots, leafs, additionalDepsToKeep);
	}

	/**
	* The set of SkyKeys kept after focusing. The actual change is done in place with the in-memory
	* graph.
	*/
	public static class FocusResult {

	private final ImmutableSet<SkyKey> rdeps;
	private final ImmutableSet<SkyKey> deps;

	private FocusResult(ImmutableSet<SkyKey> rdeps, ImmutableSet<SkyKey> deps) {
	this.rdeps = rdeps;
	this.deps = deps;
	}

	/**
	* Returns the set of SkyKeys that are in the dependencies of all roots, and rdeps from the
	* leafs. May contain transitive dependencies, in cases where certain functions use them without
	* establishing a Skyframe dependency.
	*/
	public ImmutableSet<SkyKey> getDeps() {
	return deps;
	}

	/** Returns the set of SkyKeys that are in the reverse dependencies of the leafs. */
	public ImmutableSet<SkyKey> getRdeps() {
	return rdeps;
	}
	}

	/**
	* NodeVisitor is parallelizable graph visitor that's applied transitively from leafs to the
	* roots, while marking rdeps and all direct deps of those rdeps to be kept by SkyframeFocuser.
	*/
	private class NodeVisitor implements Runnable {

	// The SkyKey that this NodeVisitor is responsible for.
	private final SkyKey key;

	// Threadsafe set of keys that depend on this key. May be modified by multiple NodeVisitors
	// concurrently.
	private final Set<SkyKey> keptRdeps;

	// Threadsafe set of keys that this key depends on. May be modified by multiple NodeVisitors
	// concurrently.
	private final Set<SkyKey> keptDeps;

	// See javadoc for the additionalDepsToKeep parameter of {@code SkyframeFocuser#focus}.
	private final Function<SkyKey, Set<SkyKey>> additionalDepsToKeep;

	protected NodeVisitor(
	SkyKey key,
	Set<SkyKey> keptRdeps,
	Set<SkyKey> keptDeps,
	Function<SkyKey, Set<SkyKey>> additionalDepsToKeep) {
	this.key = key;
	this.keptRdeps = keptRdeps;
	this.keptDeps = keptDeps;
	this.additionalDepsToKeep = additionalDepsToKeep;
	}

	@Override
	public void run() {
	InMemoryNodeEntry nodeEntry = graph.getIfPresent(key);
	if (nodeEntry == null) {
	// Throwing may be too strong if the working set is loosely defined, e.g. an entire
	// directory instead of a single file, and there are some files in the directory that
	// are not in the TC of the roots.
	//
	// TODO: b/312819241 - handle this gracefully without throwing.
	throw new IllegalStateException("nodeEntry not found for: " + key.getCanonicalName());
	}

	if (!nodeEntry.isDone()) {
	// TODO: b/312819241 - handle this gracefully without throwing.
	throw new IllegalStateException("nodeEntry not done: " + key.getCanonicalName());
	}

	for (SkyKey rdep : nodeEntry.getReverseDepsForDoneEntry()) {
	if (!keptRdeps.add(rdep)) {
	// Memoization. Already processed.
	continue;
	}

	// Queue a traversal up the graph. This will not create duplicate NodeVisitors on the
	// same rdep due to the atomic keptRdeps.add check above.
	execute(new NodeVisitor(rdep, keptRdeps, keptDeps, additionalDepsToKeep));
	}

	for (SkyKey dep : nodeEntry.getDirectDeps()) {
	if (!keptDeps.add(dep)) {
	// Memoization. Already processed.
	continue;
	}

	// This is necessary to keep the action inputs encapsulated by a NestedSet. Otherwise,
	// those inputs will be missing.
	//
	// TODO: b/312819241 - move SkyframeFocuser from build.skyframe to build.lib.skyframe so
	// that we can do the check directly without using an injected Function.
	execute(() -> keptDeps.addAll(additionalDepsToKeep.apply(dep)));
	}
	}
	}

	/** Entry point of the Skyframe garbage collection algorithm. */
	private FocusResult run(
	Set<SkyKey> roots, Set<SkyKey> leafs, Function<SkyKey, Set<SkyKey>> additionalDepsToKeep)
	throws InterruptedException {

	Set<SkyKey> keptDeps = Sets.newConcurrentHashSet();
	Set<SkyKey> keptRdeps = Sets.newConcurrentHashSet();

	// All leafs are automatically considered as rdeps.
	keptRdeps.addAll(leafs);

	// All roots are automatically considered as deps.
	//
	// Some roots are re-evaluated on every build. These roots may not be in the reverse TC
	// of leafs (working set), but may influence how the working set is evaluated (e.g. platform
	// mapping). If we remove them from the graph, those keys may be re-evaluated anyway (along with
	// their TC) on subsequent invocations, leading to wasted compute and RAM.
	// The exercise of ensuring which roots should be kept is left to the caller of
	// this function, but we ensure that all specified ones are kept here.
	keptDeps.addAll(roots);

	try (SilentCloseable c = Profiler.instance().profile("focus.mark")) {
	// Start traversal from leafs.
	for (SkyKey leaf : leafs) {
	execute(new NodeVisitor(leaf, keptRdeps, keptDeps, additionalDepsToKeep));
	}
	awaitQuiescenceWithoutShutdown(true);
	}

	// Keep the rdeps transitive closure from leafs distinct from the deps.
	keptDeps.removeAll(keptRdeps);

	eventHandler.handle(
	Event.info("Nodes in reverse transitive closure from leafs: " + keptRdeps.size()));
	eventHandler.handle(
	Event.info("Nodes in direct deps of reverse transitive closure: " + keptDeps.size()));

	try (SilentCloseable c = Profiler.instance().profile("focus.sweep_nodes")) {
	Set<SkyKey> toKeep = Sets.union(keptDeps, keptRdeps);
	graph.parallelForEach(
	inMemoryNodeEntry -> {
	SkyKey key = inMemoryNodeEntry.getKey();
	if (!toKeep.contains(key)) {
	graph.remove(key);
	}
	});

	graph.shrinkNodeMap();
	}

	AtomicLong rdepEdgesBefore = new AtomicLong();
	AtomicLong rdepEdgesAfter = new AtomicLong();

	try (SilentCloseable c = Profiler.instance().profile("focus.sweep_edges")) {
	for (SkyKey key : keptDeps) {
	execute(
	() -> {
	// TODO: b/312819241 - Consider transforming IncrementalInMemoryNodeEntry only used
	// for their immutable states to an ImmutableDoneNodeEntry or
	// NonIncrementalInMemoryNodeEntry
	// for further memory savings.
	IncrementalInMemoryNodeEntry nodeEntry =
	(IncrementalInMemoryNodeEntry) graph.getIfPresent(key);

	// No need to keep the direct deps edges of existing deps. For example:
	//
	// B
	// / \|\
	// A C \
	// \ \|
	// D
	//
	// B is the root, and A is the only leaf. We can throw out the CD edge, even
	// though both C and D are still used by B. This is because no changes are expected to
	// C and D, so it's unnecessary to maintain the edges.
	nodeEntry.directDeps = GroupedDeps.EMPTY_COMPRESSED;

	// No need to keep the rdep edges of the deps if they do not point to an rdep
	// reachable (hence, dirty-able) by the working set.
	//
	// This accounts for nearly 5% of 9+GB retained heap on a large server build.
	Collection<SkyKey> existingRdeps = nodeEntry.getReverseDepsForDoneEntry();
	rdepEdgesBefore.getAndAdd(existingRdeps.size());
	int rdepEdgesKept = 0;
	for (SkyKey rdep : existingRdeps) {
	if (keptRdeps.contains(rdep)) {
	rdepEdgesKept++;
	} else {
	nodeEntry.removeReverseDep(rdep);
	}
	}
	rdepEdgesAfter.getAndAdd(rdepEdgesKept);

	ReverseDepsUtility.consolidateData(nodeEntry);
	});
	}

	awaitQuiescence(true); // and shut down the ExecutorService.
	}

	eventHandler.handle(
	Event.info(
	String.format("Rdep edges: %s -> %s", rdepEdgesBefore.get(), rdepEdgesAfter.get())));

	return new FocusResult(ImmutableSet.copyOf(keptRdeps), ImmutableSet.copyOf(keptDeps));
	}
	}