Add functionality to MemoryProfiler to do multiple garbage collections at the end of the build in an effort to get an accurate measurement of used memory.

PiperOrigin-RevId: 187337487

6 files changed, 122 insertions, 13 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/buildtool/BuildTool.java b/src/main/java/com/google/devtools/build/lib/buildtool/BuildTool.java
index adc6b18113..c793faf1d6 100644
--- a/src/main/java/com/google/devtools/build/lib/buildtool/BuildTool.java
+++ b/src/main/java/com/google/devtools/build/lib/buildtool/BuildTool.java
@@ -300,8 +300,9 @@ public final class BuildTool {
         if (errorMessage != null) {
           throw new BuildFailedException(errorMessage);
         }
-        // Return.
+        // Will return after profiler line below.
       }
+      Profiler.instance().markPhase(ProfilePhase.FINISH);
     } catch (RuntimeException e) {
       // Print an error message for unchecked runtime exceptions. This does not concern Error
       // subclasses such as OutOfMemoryError.
diff --git a/src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java b/src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java
index 2b607d8a56..7adb3b5f93 100644
--- a/src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java
+++ b/src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java
@@ -482,8 +482,6 @@ public class ExecutionTool {
         actionContextProvider.executionPhaseEnding();
       }
 
-      Profiler.instance().markPhase(ProfilePhase.FINISH);
-
       if (buildCompleted) {
         saveActionCache(actionCache);
       }
@@ -518,7 +516,8 @@ public class ExecutionTool {
     }
   }
 
-  private void prepare(PackageRoots packageRoots) throws ExecutorInitException {
+  private void prepare(PackageRoots packageRoots)
+      throws ExecutorInitException, InterruptedException {
     Optional<ImmutableMap<PackageIdentifier, Root>> packageRootMap =
         packageRoots.getPackageRootsMap();
     if (!packageRootMap.isPresent()) {
diff --git a/src/main/java/com/google/devtools/build/lib/profiler/MemoryProfiler.java b/src/main/java/com/google/devtools/build/lib/profiler/MemoryProfiler.java
index 0be0b1c9c6..59e995a3e0 100644
--- a/src/main/java/com/google/devtools/build/lib/profiler/MemoryProfiler.java
+++ b/src/main/java/com/google/devtools/build/lib/profiler/MemoryProfiler.java
@@ -14,10 +14,18 @@
 
 package com.google.devtools.build.lib.profiler;
 
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Splitter;
+import com.google.devtools.common.options.OptionsParsingException;
 import java.io.OutputStream;
 import java.io.PrintStream;
 import java.lang.management.ManagementFactory;
+import java.lang.management.MemoryMXBean;
 import java.lang.management.MemoryUsage;
+import java.time.Duration;
+import java.util.Iterator;
+import java.util.NoSuchElementException;
+import javax.annotation.Nullable;
 
 /**
  * Blaze memory profiler.
@@ -46,10 +54,18 @@ public final class MemoryProfiler {
 
   private PrintStream memoryProfile;
   private ProfilePhase currentPhase;
+  private long heapUsedMemoryAtFinish;
+  @Nullable private MemoryProfileStableHeapParameters memoryProfileStableHeapParameters;
+
+  public synchronized void setStableMemoryParameters(
+      MemoryProfileStableHeapParameters memoryProfileStableHeapParameters) {
+    this.memoryProfileStableHeapParameters = memoryProfileStableHeapParameters;
+  }
 
   public synchronized void start(OutputStream out) {
     this.memoryProfile = (out == null) ? null : new PrintStream(out);
     this.currentPhase = ProfilePhase.INIT;
+    heapUsedMemoryAtFinish = 0;
   }
 
   public synchronized void stop() {
@@ -57,19 +73,28 @@ public final class MemoryProfiler {
       memoryProfile.close();
       memoryProfile = null;
     }
+    heapUsedMemoryAtFinish = 0;
   }
 
-  public synchronized void markPhase(ProfilePhase nextPhase) {
+  public synchronized long getHeapUsedMemoryAtFinish() {
+    return heapUsedMemoryAtFinish;
+  }
+
+  public synchronized void markPhase(ProfilePhase nextPhase) throws InterruptedException {
     if (memoryProfile != null) {
+      MemoryMXBean bean = ManagementFactory.getMemoryMXBean();
+      prepareBean(nextPhase, bean, (duration) -> Thread.sleep(duration.toMillis()));
       String name = currentPhase.description;
-      ManagementFactory.getMemoryMXBean().gc();
-      MemoryUsage memoryUsage = ManagementFactory.getMemoryMXBean().getHeapMemoryUsage();
+      MemoryUsage memoryUsage = bean.getHeapMemoryUsage();
       memoryProfile.println(name + ":heap:init:" + memoryUsage.getInit());
       memoryProfile.println(name + ":heap:used:" + memoryUsage.getUsed());
       memoryProfile.println(name + ":heap:commited:" + memoryUsage.getCommitted());
       memoryProfile.println(name + ":heap:max:" + memoryUsage.getMax());
+      if (nextPhase == ProfilePhase.FINISH) {
+        heapUsedMemoryAtFinish = memoryUsage.getUsed();
+      }
 
-      memoryUsage = ManagementFactory.getMemoryMXBean().getNonHeapMemoryUsage();
+      memoryUsage = bean.getNonHeapMemoryUsage();
       memoryProfile.println(name + ":non-heap:init:" + memoryUsage.getInit());
       memoryProfile.println(name + ":non-heap:used:" + memoryUsage.getUsed());
       memoryProfile.println(name + ":non-heap:commited:" + memoryUsage.getCommitted());
@@ -77,4 +102,71 @@ public final class MemoryProfiler {
       currentPhase = nextPhase;
     }
   }
+
+  @VisibleForTesting
+  synchronized void prepareBean(ProfilePhase nextPhase, MemoryMXBean bean, Sleeper sleeper)
+      throws InterruptedException {
+    bean.gc();
+    if (nextPhase == ProfilePhase.FINISH && memoryProfileStableHeapParameters != null) {
+      for (int i = 1; i < memoryProfileStableHeapParameters.numTimesToDoGc; i++) {
+        sleeper.sleep(memoryProfileStableHeapParameters.timeToSleepBetweenGcs);
+        bean.gc();
+      }
+    }
+  }
+
+  /**
+   * Parameters that control how {@code MemoryProfiler} tries to get a stable heap at the end of the
+   * build.
+   */
+  public static class MemoryProfileStableHeapParameters {
+    private final int numTimesToDoGc;
+    private final Duration timeToSleepBetweenGcs;
+
+    private MemoryProfileStableHeapParameters(int numTimesToDoGc, Duration timeToSleepBetweenGcs) {
+      this.numTimesToDoGc = numTimesToDoGc;
+      this.timeToSleepBetweenGcs = timeToSleepBetweenGcs;
+    }
+
+    /** Converter for {@code MemoryProfileStableHeapParameters} option. */
+    public static class Converter
+        implements com.google.devtools.common.options.Converter<MemoryProfileStableHeapParameters> {
+      private static final Splitter SPLITTER = Splitter.on(',');
+
+      @Override
+      public MemoryProfileStableHeapParameters convert(String input)
+          throws OptionsParsingException {
+        Iterator<String> values = SPLITTER.split(input).iterator();
+        try {
+          int numTimesToDoGc = Integer.parseInt(values.next());
+          int numSecondsToSleepBetweenGcs = Integer.parseInt(values.next());
+          if (values.hasNext()) {
+            throw new OptionsParsingException("Expected exactly 2 comma-separated integer values");
+          }
+          if (numTimesToDoGc <= 0) {
+            throw new OptionsParsingException("Number of times to GC must be positive");
+          }
+          if (numSecondsToSleepBetweenGcs < 0) {
+            throw new OptionsParsingException(
+                "Number of seconds to sleep between GC's must be positive");
+          }
+          return new MemoryProfileStableHeapParameters(
+              numTimesToDoGc, Duration.ofSeconds(numSecondsToSleepBetweenGcs));
+        } catch (NumberFormatException | NoSuchElementException nfe) {
+          throw new OptionsParsingException(
+              "Expected exactly 2 comma-separated integer values", nfe);
+        }
+      }
+
+      @Override
+      public String getTypeDescription() {
+        return "two integers, separated by a comma";
+      }
+    }
+  }
+
+  @VisibleForTesting
+  interface Sleeper {
+    void sleep(Duration duration) throws InterruptedException;
+  }
 }
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
index 0b126555d3..d453967041 100644
--- a/src/main/java/com/google/devtools/build/lib/profiler/Profiler.java
+++ b/src/main/java/com/google/devtools/build/lib/profiler/Profiler.java
@@ -886,10 +886,8 @@ public final class Profiler {
     }
   }
 
-  /**
-   * Convenience method to log phase marker tasks.
-   */
-  public void markPhase(ProfilePhase phase) {
+  /** Convenience method to log phase marker tasks. */
+  public void markPhase(ProfilePhase phase) throws InterruptedException {
     MemoryProfiler.instance().markPhase(phase);
     if (isActive() && isProfiling(ProfilerTask.PHASE)) {
       Preconditions.checkState(taskStack.isEmpty(), "Phase tasks must not be nested");
diff --git a/src/main/java/com/google/devtools/build/lib/runtime/BlazeRuntime.java b/src/main/java/com/google/devtools/build/lib/runtime/BlazeRuntime.java
index df1692fdd5..6932d9ea18 100644
--- a/src/main/java/com/google/devtools/build/lib/runtime/BlazeRuntime.java
+++ b/src/main/java/com/google/devtools/build/lib/runtime/BlazeRuntime.java
@@ -407,6 +407,8 @@ public final class BlazeRuntime {
 
     if (options.memoryProfilePath != null) {
       Path memoryProfilePath = env.getWorkingDirectory().getRelative(options.memoryProfilePath);
+      MemoryProfiler.instance()
+          .setStableMemoryParameters(options.memoryProfileStableHeapParameters);
       try {
         MemoryProfiler.instance().start(memoryProfilePath.getOutputStream());
       } catch (IOException e) {
diff --git a/src/main/java/com/google/devtools/build/lib/runtime/CommonCommandOptions.java b/src/main/java/com/google/devtools/build/lib/runtime/CommonCommandOptions.java
index 238eb225d6..ea9c7fa057 100644
--- a/src/main/java/com/google/devtools/build/lib/runtime/CommonCommandOptions.java
+++ b/src/main/java/com/google/devtools/build/lib/runtime/CommonCommandOptions.java
@@ -15,6 +15,7 @@ package com.google.devtools.build.lib.runtime;
 
 import static com.google.common.base.Strings.isNullOrEmpty;
 
+import com.google.devtools.build.lib.profiler.MemoryProfiler.MemoryProfileStableHeapParameters;
 import com.google.devtools.build.lib.runtime.CommandLineEvent.ToolCommandLineEvent;
 import com.google.devtools.build.lib.util.OptionsUtils;
 import com.google.devtools.build.lib.vfs.PathFragment;
@@ -229,11 +230,26 @@ public class CommonCommandOptions extends OptionsBase {
     documentationCategory = OptionDocumentationCategory.UNDOCUMENTED,
     effectTags = {OptionEffectTag.AFFECTS_OUTPUTS, OptionEffectTag.BAZEL_MONITORING},
     converter = OptionsUtils.PathFragmentConverter.class,
-    help = "If set, write memory usage data to the specified file at phase ends."
+    help =
+        "If set, write memory usage data to the specified file at phase ends and stable heap to"
+            + " master log at end of build."
   )
   public PathFragment memoryProfilePath;
 
   @Option(
+    name = "memory_profile_stable_heap_parameters",
+    defaultValue = "1,0",
+    documentationCategory = OptionDocumentationCategory.LOGGING,
+    effectTags = {OptionEffectTag.BAZEL_MONITORING},
+    converter = MemoryProfileStableHeapParameters.Converter.class,
+    help =
+        "Tune memory profile's computation of stable heap at end of build. Should be two integers "
+            + "separated by a comma. First parameter is the number of GCs to perform. Second "
+            + "parameter is the number of seconds to wait between GCs."
+  )
+  public MemoryProfileStableHeapParameters memoryProfileStableHeapParameters;
+
+  @Option(
     name = "experimental_oom_more_eagerly_threshold",
     defaultValue = "100",
     documentationCategory = OptionDocumentationCategory.EXECUTION_STRATEGY,
@@ -362,4 +378,5 @@ public class CommonCommandOptions extends OptionsBase {
             + "or the bad combination should be checked for programmatically."
   )
   public List<String> deprecationWarnings;
+
 }