Makes dumping `starlark_memory` adopt `InstrumentationOutput`
PiperOrigin-RevId: 696656280
Change-Id: I9913d561674dd5547b13333a536167c0e4601ab0
diff --git a/src/main/java/com/google/devtools/build/lib/profiler/memory/AllocationTracker.java b/src/main/java/com/google/devtools/build/lib/profiler/memory/AllocationTracker.java
index 3611cee..b21af90 100644
--- a/src/main/java/com/google/devtools/build/lib/profiler/memory/AllocationTracker.java
+++ b/src/main/java/com/google/devtools/build/lib/profiler/memory/AllocationTracker.java
@@ -31,8 +31,8 @@
import com.google.perftools.profiles.ProfileProto.Profile;
import com.google.perftools.profiles.ProfileProto.Sample;
import com.google.perftools.profiles.ProfileProto.ValueType;
-import java.io.FileOutputStream;
import java.io.IOException;
+import java.io.OutputStream;
import java.time.Instant;
import java.util.HashMap;
import java.util.Map;
@@ -281,14 +281,14 @@
}
/** Dumps all Starlark analysis time allocations to a pprof-compatible file. */
- public void dumpStarlarkAllocations(String path) throws IOException {
+ public void dumpStarlarkAllocations(OutputStream outputStream) throws IOException {
// Make sure we don't track our own allocations
enabled = false;
System.gc();
Profile profile = buildMemoryProfile();
- try (GZIPOutputStream outputStream = new GZIPOutputStream(new FileOutputStream(path))) {
- profile.writeTo(outputStream);
- outputStream.finish();
+ try (GZIPOutputStream gzipOutputStream = new GZIPOutputStream(outputStream)) {
+ profile.writeTo(gzipOutputStream);
+ gzipOutputStream.finish();
}
enabled = true;
}
diff --git a/src/main/java/com/google/devtools/build/lib/runtime/commands/DumpCommand.java b/src/main/java/com/google/devtools/build/lib/runtime/commands/DumpCommand.java
index 5b1bf8d..073a4c0 100644
--- a/src/main/java/com/google/devtools/build/lib/runtime/commands/DumpCommand.java
+++ b/src/main/java/com/google/devtools/build/lib/runtime/commands/DumpCommand.java
@@ -39,6 +39,8 @@
import com.google.devtools.build.lib.runtime.BlazeWorkspace;
import com.google.devtools.build.lib.runtime.Command;
import com.google.devtools.build.lib.runtime.CommandEnvironment;
+import com.google.devtools.build.lib.runtime.InstrumentationOutput;
+import com.google.devtools.build.lib.runtime.InstrumentationOutputFactory.DestinationRelativeTo;
import com.google.devtools.build.lib.server.FailureDetails;
import com.google.devtools.build.lib.server.FailureDetails.DumpCommand.Code;
import com.google.devtools.build.lib.server.FailureDetails.FailureDetail;
@@ -52,6 +54,7 @@
import com.google.devtools.build.lib.util.MemoryAccountant.Stats;
import com.google.devtools.build.lib.util.RegexFilter;
import com.google.devtools.build.lib.util.RegexFilter.RegexFilterConverter;
+import com.google.devtools.build.lib.vfs.PathFragment;
import com.google.devtools.build.skyframe.InMemoryGraph;
import com.google.devtools.build.skyframe.MemoizingEvaluator;
import com.google.devtools.build.skyframe.NodeEntry;
@@ -384,7 +387,19 @@
if (dumpOptions.starlarkMemory != null) {
try {
- dumpStarlarkHeap(env.getBlazeWorkspace(), dumpOptions.starlarkMemory, out);
+ InstrumentationOutput starlarkHeapOutput =
+ runtime
+ .getInstrumentationOutputFactory()
+ .createInstrumentationOutput(
+ /* name= */ "starlark_heap",
+ PathFragment.create(dumpOptions.starlarkMemory),
+ DestinationRelativeTo.WORKSPACE_OR_HOME,
+ env,
+ env.getReporter(),
+ /* append= */ null,
+ /* internal= */ null);
+ dumpStarlarkHeap(
+ env.getBlazeWorkspace(), starlarkHeapOutput, dumpOptions.starlarkMemory, out);
} catch (IOException e) {
String message = "Could not dump Starlark memory";
env.getReporter().error(null, message, e);
@@ -672,7 +687,11 @@
return Optional.empty();
}
- private static void dumpStarlarkHeap(BlazeWorkspace workspace, String path, PrintStream out)
+ private static void dumpStarlarkHeap(
+ BlazeWorkspace workspace,
+ InstrumentationOutput starlarkHeapOutput,
+ String path,
+ PrintStream out)
throws IOException {
AllocationTracker allocationTracker = workspace.getAllocationTracker();
if (allocationTracker == null) {
@@ -683,7 +702,10 @@
return;
}
out.println("Dumping Starlark heap to: " + path);
- allocationTracker.dumpStarlarkAllocations(path);
+
+ // OutputStream is expected to be closed when allocationTracker.dumpStarlarkAllocations()
+ // returns.
+ allocationTracker.dumpStarlarkAllocations(starlarkHeapOutput.createOutputStream());
}
static BlazeCommandResult createFailureResult(String message, Code detailedCode) {
