Introduce `InstrumentationOutput` and apply it to profile `InstrumentationOutput` aims at providing a unified way to handle writing and publishing instrumentation outputs. PiperOrigin-RevId: 663316138 Change-Id: Ie5d80acd1843a94045307442ed00ffd6a9f98105
diff --git a/src/main/java/com/google/devtools/build/lib/buildtool/buildevent/ProfilerStartedEvent.java b/src/main/java/com/google/devtools/build/lib/buildtool/buildevent/ProfilerStartedEvent.java index fcb2f94..6c89b12 100644 --- a/src/main/java/com/google/devtools/build/lib/buildtool/buildevent/ProfilerStartedEvent.java +++ b/src/main/java/com/google/devtools/build/lib/buildtool/buildevent/ProfilerStartedEvent.java
@@ -13,46 +13,20 @@ // limitations under the License. package com.google.devtools.build.lib.buildtool.buildevent; -import com.google.devtools.build.lib.buildeventstream.BuildEventArtifactUploader.UploadContext; import com.google.devtools.build.lib.events.ExtendedEventHandler; -import com.google.devtools.build.lib.profiler.Profiler; -import com.google.devtools.build.lib.vfs.Path; +import com.google.devtools.build.lib.runtime.InstrumentationOutput; import javax.annotation.Nullable; /** This event is fired when the profiler is started. */ public class ProfilerStartedEvent implements ExtendedEventHandler.Postable { - @Nullable private final Path profilePath; - @Nullable private final UploadContext streamingContext; - private final Profiler.Format format; + @Nullable private final InstrumentationOutput profile; - public ProfilerStartedEvent( - @Nullable Path profilePath, - @Nullable UploadContext streamingContext, - Profiler.Format format) { - this.profilePath = profilePath; - this.streamingContext = streamingContext; - this.format = format; + public ProfilerStartedEvent(@Nullable InstrumentationOutput profile) { + this.profile = profile; } @Nullable - public Path getProfilePath() { - return profilePath; - } - - @Nullable - public UploadContext getStreamingContext() { - return streamingContext; - } - - public String getName() { - switch (format) { - case JSON_TRACE_FILE_FORMAT -> { - return "command.profile.json"; - } - case JSON_TRACE_FILE_COMPRESSED_FORMAT -> { - return "command.profile.gz"; - } - } - throw new UnsupportedOperationException(); + public InstrumentationOutput getProfile() { + return profile; } }
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 ec89a88..98af82e 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
@@ -43,9 +43,7 @@ import com.google.devtools.build.lib.bugreport.BugReporter; import com.google.devtools.build.lib.bugreport.Crash; import com.google.devtools.build.lib.bugreport.CrashContext; -import com.google.devtools.build.lib.buildeventstream.BuildEvent.LocalFile.LocalFileType; import com.google.devtools.build.lib.buildeventstream.BuildEventArtifactUploader; -import com.google.devtools.build.lib.buildeventstream.BuildEventArtifactUploader.UploadContext; import com.google.devtools.build.lib.buildeventstream.BuildEventProtocolOptions; import com.google.devtools.build.lib.buildtool.CommandPrecompleteEvent; import com.google.devtools.build.lib.buildtool.buildevent.ProfilerStartedEvent; @@ -324,27 +322,38 @@ ImmutableSet.Builder<ProfilerTask> profiledTasksBuilder = ImmutableSet.builder(); Profiler.Format format = Format.JSON_TRACE_FILE_FORMAT; Path profilePath = null; - UploadContext streamingContext = null; + InstrumentationOutput profile = null; try { if (tracerEnabled) { if (options.profilePath == null) { + String profileName = "command.profile.gz"; format = Format.JSON_TRACE_FILE_COMPRESSED_FORMAT; if (bepOptions != null && bepOptions.streamingLogFileUploads) { BuildEventArtifactUploader buildEventArtifactUploader = newUploader(env, bepOptions.buildEventUploadStrategy); - streamingContext = buildEventArtifactUploader.startUpload(LocalFileType.LOG, null); - out = streamingContext.getOutputStream(); + profile = + new BuildEventArtifactInstrumentationOutput( + profileName, buildEventArtifactUploader); + out = profile.createOutputStream(); } else { - profilePath = workspace.getOutputBase().getRelative("command.profile.gz"); - out = profilePath.getOutputStream(); + profilePath = workspace.getOutputBase().getRelative(profileName); + profile = new LocalInstrumentationOutput(profileName, profilePath); + out = profile.createOutputStream(); } } else { format = options.profilePath.toString().endsWith(".gz") ? Format.JSON_TRACE_FILE_COMPRESSED_FORMAT : Format.JSON_TRACE_FILE_FORMAT; + String profileName = + (format == Format.JSON_TRACE_FILE_COMPRESSED_FORMAT) + ? "command.profile.gz" + : "command.profile.json"; profilePath = workspace.getWorkspace().getRelative(options.profilePath); - out = profilePath.getOutputStream(/* append= */ false, /* internal= */ true); + profile = new LocalInstrumentationOutput(profileName, profilePath); + out = + ((LocalInstrumentationOutput) profile) + .createOutputStream(/* append= */ false, /* internal= */ true); } for (ProfilerTask profilerTask : ProfilerTask.values()) { if (!profilerTask.isVfs() @@ -445,7 +454,7 @@ } catch (IOException e) { eventHandler.handle(Event.error("Error while creating profile file: " + e.getMessage())); } - return new ProfilerStartedEvent(profilePath, streamingContext, format); + return new ProfilerStartedEvent(profile); } public FileSystem getFileSystem() {
diff --git a/src/main/java/com/google/devtools/build/lib/runtime/BuildEventArtifactInstrumentationOutput.java b/src/main/java/com/google/devtools/build/lib/runtime/BuildEventArtifactInstrumentationOutput.java new file mode 100644 index 0000000..dad4993 --- /dev/null +++ b/src/main/java/com/google/devtools/build/lib/runtime/BuildEventArtifactInstrumentationOutput.java
@@ -0,0 +1,51 @@ +// 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.runtime; + +import static com.google.common.base.Preconditions.checkNotNull; + +import com.google.devtools.build.lib.buildeventstream.BuildEvent.LocalFile.LocalFileType; +import com.google.devtools.build.lib.buildeventstream.BuildEventArtifactUploader; +import com.google.devtools.build.lib.buildeventstream.BuildEventArtifactUploader.UploadContext; +import com.google.devtools.build.lib.buildtool.BuildResult.BuildToolLogCollection; +import java.io.OutputStream; +import javax.annotation.Nullable; + +/** + * Used when instrumentation output should be treated as a build event artifact so that it will be + * uploaded to a specified location. + */ +final class BuildEventArtifactInstrumentationOutput implements InstrumentationOutput { + private final String name; + private final BuildEventArtifactUploader buildEventArtifactUploader; + @Nullable private UploadContext uploadContext; + + public BuildEventArtifactInstrumentationOutput( + String name, BuildEventArtifactUploader buildEventArtifactUploader) { + this.name = checkNotNull(name); + this.buildEventArtifactUploader = checkNotNull(buildEventArtifactUploader); + } + + @Override + public void publish(BuildToolLogCollection buildToolLogCollection) { + checkNotNull(uploadContext, "Cannot publish to buildToolLogCollection if upload never starts."); + buildToolLogCollection.addUriFuture(name, uploadContext.uriFuture()); + } + + @Override + public OutputStream createOutputStream() { + uploadContext = buildEventArtifactUploader.startUpload(LocalFileType.LOG, null); + return uploadContext.getOutputStream(); + } +}
diff --git a/src/main/java/com/google/devtools/build/lib/runtime/BuildSummaryStatsModule.java b/src/main/java/com/google/devtools/build/lib/runtime/BuildSummaryStatsModule.java index 70fe506..8a865f7 100644 --- a/src/main/java/com/google/devtools/build/lib/runtime/BuildSummaryStatsModule.java +++ b/src/main/java/com/google/devtools/build/lib/runtime/BuildSummaryStatsModule.java
@@ -179,16 +179,13 @@ } } } - if (profileEvent != null && profileEvent.getProfilePath() != null) { + if (profileEvent != null && profileEvent.getProfile() != null) { // This leads to missing the afterCommand profiles of the other modules in the profile. // Since the BEP currently shuts down at the BuildCompleteEvent, we cannot just move posting // the BuildToolLogs to afterCommand of this module. try { Profiler.instance().stop(); - event - .getResult() - .getBuildToolLogCollection() - .addLocalFile(profileEvent.getName(), profileEvent.getProfilePath()); + profileEvent.getProfile().publish(event.getResult().getBuildToolLogCollection()); } catch (IOException e) { reporter.handle(Event.error("Error while writing profile file: " + e.getMessage())); }
diff --git a/src/main/java/com/google/devtools/build/lib/runtime/InstrumentationOutput.java b/src/main/java/com/google/devtools/build/lib/runtime/InstrumentationOutput.java new file mode 100644 index 0000000..f5130e6 --- /dev/null +++ b/src/main/java/com/google/devtools/build/lib/runtime/InstrumentationOutput.java
@@ -0,0 +1,28 @@ +// 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.runtime; + +import com.google.devtools.build.lib.buildtool.BuildResult.BuildToolLogCollection; +import java.io.IOException; +import java.io.OutputStream; + +/** Stores and publishes the instrumentation output information. */ +public interface InstrumentationOutput { + + /** Creates the {@link OutputStream} for instrumentation output writes. */ + OutputStream createOutputStream() throws IOException; + + /** Publishes instrumentation output information to the {@link BuildToolLogCollection}. */ + void publish(BuildToolLogCollection buildToolLogCollection); +}
diff --git a/src/main/java/com/google/devtools/build/lib/runtime/LocalInstrumentationOutput.java b/src/main/java/com/google/devtools/build/lib/runtime/LocalInstrumentationOutput.java new file mode 100644 index 0000000..aedb96a --- /dev/null +++ b/src/main/java/com/google/devtools/build/lib/runtime/LocalInstrumentationOutput.java
@@ -0,0 +1,44 @@ +// 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.runtime; + +import com.google.devtools.build.lib.buildtool.BuildResult.BuildToolLogCollection; +import com.google.devtools.build.lib.vfs.Path; +import java.io.IOException; +import java.io.OutputStream; + +/** Used when instrumentation output is written to a local file. */ +final class LocalInstrumentationOutput implements InstrumentationOutput { + private final Path path; + private final String name; + + LocalInstrumentationOutput(String name, Path path) { + this.name = name; + this.path = path; + } + + @Override + public void publish(BuildToolLogCollection buildToolLogCollection) { + buildToolLogCollection.addLocalFile(name, path); + } + + @Override + public OutputStream createOutputStream() throws IOException { + return path.getOutputStream(); + } + + public OutputStream createOutputStream(boolean append, boolean internal) throws IOException { + return path.getOutputStream(append, internal); + } +}
diff --git a/src/test/java/com/google/devtools/build/lib/buildtool/buildevent/BUILD b/src/test/java/com/google/devtools/build/lib/buildtool/buildevent/BUILD deleted file mode 100644 index e3a3672..0000000 --- a/src/test/java/com/google/devtools/build/lib/buildtool/buildevent/BUILD +++ /dev/null
@@ -1,29 +0,0 @@ -load("@rules_java//java:defs.bzl", "java_test") - -package( - default_applicable_licenses = ["//:license"], - default_testonly = 1, - default_visibility = ["//src:__subpackages__"], -) - -filegroup( - name = "srcs", - testonly = 0, - srcs = glob(["**"]), - visibility = ["//src:__subpackages__"], -) - -java_test( - name = "ProfilerStartedEventTest", - srcs = ["ProfilerStartedEventTest.java"], - deps = [ - "//src/main/java/com/google/devtools/build/lib:runtime", - "//src/main/java/com/google/devtools/build/lib/buildeventstream", - "//src/main/java/com/google/devtools/build/lib/profiler", - "//src/main/java/com/google/devtools/build/lib/vfs", - "//src/main/java/com/google/devtools/build/lib/vfs/inmemoryfs", - "//third_party:junit4", - "//third_party:mockito", - "//third_party:truth", - ], -)
diff --git a/src/test/java/com/google/devtools/build/lib/buildtool/buildevent/ProfilerStartedEventTest.java b/src/test/java/com/google/devtools/build/lib/buildtool/buildevent/ProfilerStartedEventTest.java deleted file mode 100644 index 628fb3b..0000000 --- a/src/test/java/com/google/devtools/build/lib/buildtool/buildevent/ProfilerStartedEventTest.java +++ /dev/null
@@ -1,57 +0,0 @@ -// 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.buildtool.buildevent; - -import static com.google.common.truth.Truth.assertThat; -import static org.mockito.Mockito.mock; - -import com.google.devtools.build.lib.buildeventstream.BuildEventArtifactUploader.UploadContext; -import com.google.devtools.build.lib.profiler.Profiler.Format; -import com.google.devtools.build.lib.vfs.DigestHashFunction; -import com.google.devtools.build.lib.vfs.FileSystem; -import com.google.devtools.build.lib.vfs.Path; -import com.google.devtools.build.lib.vfs.inmemoryfs.InMemoryFileSystem; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; - -/** Tests for {@link ProfilerStartedEvent} */ -@RunWith(JUnit4.class) -public class ProfilerStartedEventTest { - @Test - public void testLocalJsonProfiler() { - FileSystem fs = new InMemoryFileSystem(DigestHashFunction.SHA256); - Path profilePath = fs.getPath("/tmp/profile"); - ProfilerStartedEvent jsonProfilerStartedEvent = - new ProfilerStartedEvent( - profilePath, /* streamingContext= */ null, Format.JSON_TRACE_FILE_FORMAT); - - assertThat(jsonProfilerStartedEvent.getProfilePath()).isSameInstanceAs(profilePath); - assertThat(jsonProfilerStartedEvent.getStreamingContext()).isNull(); - assertThat(jsonProfilerStartedEvent.getName()).isEqualTo("command.profile.json"); - } - - @Test - public void testCompressedProfilerWithBepUploadContext() { - UploadContext streamingContext = mock(UploadContext.class); - ProfilerStartedEvent compressedProfilerStartedEvent = - new ProfilerStartedEvent( - /* profilePath= */ null, streamingContext, Format.JSON_TRACE_FILE_COMPRESSED_FORMAT); - - assertThat(compressedProfilerStartedEvent.getProfilePath()).isNull(); - assertThat(compressedProfilerStartedEvent.getStreamingContext()) - .isSameInstanceAs(streamingContext); - assertThat(compressedProfilerStartedEvent.getName()).isEqualTo("command.profile.gz"); - } -}
diff --git a/src/test/java/com/google/devtools/build/lib/runtime/BuildEventArtifactInstrumentationOutputTest.java b/src/test/java/com/google/devtools/build/lib/runtime/BuildEventArtifactInstrumentationOutputTest.java new file mode 100644 index 0000000..a5c7555 --- /dev/null +++ b/src/test/java/com/google/devtools/build/lib/runtime/BuildEventArtifactInstrumentationOutputTest.java
@@ -0,0 +1,86 @@ +// 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.runtime; + +import static com.google.common.truth.Truth.assertThat; +import static com.google.common.util.concurrent.Futures.immediateFuture; +import static org.junit.Assert.assertThrows; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.google.common.util.concurrent.ListenableFuture; +import com.google.devtools.build.lib.buildeventstream.BuildEvent.LocalFile.LocalFileType; +import com.google.devtools.build.lib.buildeventstream.BuildEventArtifactUploader; +import com.google.devtools.build.lib.buildeventstream.BuildEventArtifactUploader.UploadContext; +import com.google.devtools.build.lib.buildtool.BuildResult.BuildToolLogCollection; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.OutputStream; +import java.util.concurrent.ExecutionException; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class BuildEventArtifactInstrumentationOutputTest { + @Test + public void testBepInstrumentation_cannotPublishIfUploadNeverStarts() { + BuildEventArtifactUploader fakeBuildEventArtifactUploader = + mock(BuildEventArtifactUploader.class); + InstrumentationOutput bepInstrumentationOutput = + new BuildEventArtifactInstrumentationOutput("bep", fakeBuildEventArtifactUploader); + + BuildToolLogCollection buildToolLogCollection = new BuildToolLogCollection(); + assertThrows( + NullPointerException.class, () -> bepInstrumentationOutput.publish(buildToolLogCollection)); + } + + @Test + public void testBepInstrumentation_publishNameAndUriFuture() + throws ExecutionException, InterruptedException, IOException { + UploadContext fakeUploadLoadContext = + new UploadContext() { + @Override + public OutputStream getOutputStream() { + return new ByteArrayOutputStream(); + } + + @Override + public ListenableFuture<String> uriFuture() { + return immediateFuture("uri/abc12345"); + } + }; + BuildEventArtifactUploader fakeBuildEventArtifactUploader = + mock(BuildEventArtifactUploader.class); + when(fakeBuildEventArtifactUploader.startUpload(LocalFileType.LOG, null)) + .thenReturn(fakeUploadLoadContext); + + InstrumentationOutput bepInstrumentationOutput = + new BuildEventArtifactInstrumentationOutput("bep", fakeBuildEventArtifactUploader); + // Create the OutputStream will enforce fakeBuildEventArtifactUploader to create the + // uploadContext. + var unused = bepInstrumentationOutput.createOutputStream(); + assertThat(bepInstrumentationOutput) + .isInstanceOf(BuildEventArtifactInstrumentationOutput.class); + + BuildToolLogCollection buildToolLogCollection = new BuildToolLogCollection(); + bepInstrumentationOutput.publish(buildToolLogCollection); + buildToolLogCollection.freeze(); + + assertThat(buildToolLogCollection.toEvent().remoteUploads()).hasSize(1); + ListenableFuture<String> soleRemoteUploadUri = + buildToolLogCollection.toEvent().remoteUploads().get(0); + assertThat(soleRemoteUploadUri.get()).isEqualTo("uri/abc12345"); + } +}
diff --git a/src/test/java/com/google/devtools/build/lib/runtime/LocalInstrumentationOutputTest.java b/src/test/java/com/google/devtools/build/lib/runtime/LocalInstrumentationOutputTest.java new file mode 100644 index 0000000..e5ae19a8 --- /dev/null +++ b/src/test/java/com/google/devtools/build/lib/runtime/LocalInstrumentationOutputTest.java
@@ -0,0 +1,51 @@ +// 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.runtime; + +import static com.google.common.truth.Truth.assertThat; + +import com.google.devtools.build.lib.buildeventstream.BuildEvent.LocalFile; +import com.google.devtools.build.lib.buildeventstream.BuildEvent.LocalFile.LocalFileType; +import com.google.devtools.build.lib.buildeventstream.BuildToolLogs.LogFileEntry; +import com.google.devtools.build.lib.buildtool.BuildResult.BuildToolLogCollection; +import com.google.devtools.build.lib.vfs.DigestHashFunction; +import com.google.devtools.build.lib.vfs.FileSystem; +import com.google.devtools.build.lib.vfs.Path; +import com.google.devtools.build.lib.vfs.inmemoryfs.InMemoryFileSystem; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public final class LocalInstrumentationOutputTest { + @Test + public void testLocalInstrumentation_publishNameAndPath() { + FileSystem fs = new InMemoryFileSystem(DigestHashFunction.SHA256); + Path path = fs.getPath("/file"); + InstrumentationOutput localInstrumentationOutput = + new LocalInstrumentationOutput("local", path); + + assertThat(localInstrumentationOutput).isInstanceOf(LocalInstrumentationOutput.class); + BuildToolLogCollection buildToolLogCollection = new BuildToolLogCollection(); + localInstrumentationOutput.publish(buildToolLogCollection); + buildToolLogCollection.freeze(); + + assertThat(buildToolLogCollection.getLocalFiles()) + .containsExactly( + new LogFileEntry( + "local", + new LocalFile( + path, LocalFileType.LOG, /* artifact= */ null, /* artifactMetadata= */ null))); + } +}