From ad10633291532908db5ff83e8a446013b70e67b8 Mon Sep 17 00:00:00 2001 From: Keith Lea <79670223+keithl-stripe@users.noreply.github.com> Date: Tue, 26 Nov 2024 23:50:13 -0800 Subject: [PATCH] Add `bazel query --output_file` option, which writes query results directly to a file MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This is a proposed fix for https://github.com/bazelbuild/bazel/issues/24293 This speeds up a fully warm `bazel query ...` by 23.7%, reducing wall time from 1m49s to 1m23s ``` $ time bazel query '...' --output=streamed_proto > queryoutput4.streamedproto real 1m48.768s user 0m27.410s sys 0m19.646s $ time bazel query '...' --output=streamed_proto --output_file=queryoutput5.streamedproto real 1m22.920s user 0m0.045s sys 0m0.016s ``` _💁‍♂️ Note: when combined with https://github.com/bazelbuild/bazel/pull/24305, total wall time is 37s, an overall reduction of 66%._ Closes #24298. PiperOrigin-RevId: 700583890 Change-Id: Ic13f0611aca60c2ce8641e72a0fcfc330f13c803 --- scripts/bash_completion_test.sh | 2 +- .../java/com/google/devtools/build/lib/BUILD | 1 + .../build/lib/buildtool/AqueryProcessor.java | 9 ++- .../build/lib/buildtool/CqueryProcessor.java | 6 ++ .../buildtool/PostAnalysisQueryProcessor.java | 5 +- .../lib/query2/common/CommonQueryOptions.java | 15 +++++ .../build/lib/runtime/QueryRuntimeHelper.java | 62 ++++++++++++++++++- .../QueryEnvironmentBasedCommand.java | 2 +- .../shell/integration/bazel_query_test.sh | 18 +++++- 9 files changed, 111 insertions(+), 9 deletions(-) diff --git a/scripts/bash_completion_test.sh b/scripts/bash_completion_test.sh index bc243da2aac954..1239976a09a337 100755 --- a/scripts/bash_completion_test.sh +++ b/scripts/bash_completion_test.sh @@ -452,7 +452,7 @@ test_build_options() { test_query_options() { assert_expansion 'query --out' \ - 'query --output=' + 'query --output' # Basic label expansion works for query, too. make_packages diff --git a/src/main/java/com/google/devtools/build/lib/BUILD b/src/main/java/com/google/devtools/build/lib/BUILD index 8121ecbad79574..52cc17a363d815 100644 --- a/src/main/java/com/google/devtools/build/lib/BUILD +++ b/src/main/java/com/google/devtools/build/lib/BUILD @@ -404,6 +404,7 @@ java_library( "//src/main/java/com/google/devtools/build/lib/query2", "//src/main/java/com/google/devtools/build/lib/query2:aquery-utils", "//src/main/java/com/google/devtools/build/lib/query2/common:cquery-node", + "//src/main/java/com/google/devtools/build/lib/query2/common:options", "//src/main/java/com/google/devtools/build/lib/query2/engine", "//src/main/java/com/google/devtools/build/lib/query2/query/output", "//src/main/java/com/google/devtools/build/lib/rules:repository/repository_function", diff --git a/src/main/java/com/google/devtools/build/lib/buildtool/AqueryProcessor.java b/src/main/java/com/google/devtools/build/lib/buildtool/AqueryProcessor.java index be2787e860ca38..f5ef0b124512d3 100644 --- a/src/main/java/com/google/devtools/build/lib/buildtool/AqueryProcessor.java +++ b/src/main/java/com/google/devtools/build/lib/buildtool/AqueryProcessor.java @@ -67,11 +67,16 @@ public AqueryProcessor( actionFilters = buildActionFilters(queryExpression); } + @Override + protected AqueryOptions getQueryOptions(CommandEnvironment env) { + return env.getOptions().getOptions(AqueryOptions.class); + } + /** Outputs the current action graph from Skyframe. */ public BlazeCommandResult dumpActionGraphFromSkyframe(CommandEnvironment env) { + AqueryOptions aqueryOptions = getQueryOptions(env); try (QueryRuntimeHelper queryRuntimeHelper = - env.getRuntime().getQueryRuntimeHelperFactory().create(env)) { - AqueryOptions aqueryOptions = env.getOptions().getOptions(AqueryOptions.class); + env.getRuntime().getQueryRuntimeHelperFactory().create(env, aqueryOptions)) { PrintStream printStream = queryRuntimeHelper.getOutputStreamForQueryOutput() == null diff --git a/src/main/java/com/google/devtools/build/lib/buildtool/CqueryProcessor.java b/src/main/java/com/google/devtools/build/lib/buildtool/CqueryProcessor.java index 2ad4bc43e41242..70b15b23ed605c 100644 --- a/src/main/java/com/google/devtools/build/lib/buildtool/CqueryProcessor.java +++ b/src/main/java/com/google/devtools/build/lib/buildtool/CqueryProcessor.java @@ -19,6 +19,7 @@ import com.google.devtools.build.lib.cmdline.TargetPattern; import com.google.devtools.build.lib.packages.semantics.BuildLanguageOptions; import com.google.devtools.build.lib.query2.PostAnalysisQueryEnvironment.TopLevelConfigurations; +import com.google.devtools.build.lib.query2.common.CommonQueryOptions; import com.google.devtools.build.lib.query2.common.CqueryNode; import com.google.devtools.build.lib.query2.cquery.ConfiguredTargetQueryEnvironment; import com.google.devtools.build.lib.query2.cquery.CqueryOptions; @@ -36,6 +37,11 @@ public CqueryProcessor( super(queryExpression, mainRepoTargetParser); } + @Override + protected CommonQueryOptions getQueryOptions(CommandEnvironment env) { + return env.getOptions().getOptions(CqueryOptions.class); + } + @Override protected ConfiguredTargetQueryEnvironment getQueryEnvironment( BuildRequest request, diff --git a/src/main/java/com/google/devtools/build/lib/buildtool/PostAnalysisQueryProcessor.java b/src/main/java/com/google/devtools/build/lib/buildtool/PostAnalysisQueryProcessor.java index ec896274df845d..048dc2966ec2ca 100644 --- a/src/main/java/com/google/devtools/build/lib/buildtool/PostAnalysisQueryProcessor.java +++ b/src/main/java/com/google/devtools/build/lib/buildtool/PostAnalysisQueryProcessor.java @@ -26,6 +26,7 @@ import com.google.devtools.build.lib.query2.NamedThreadSafeOutputFormatterCallback; import com.google.devtools.build.lib.query2.PostAnalysisQueryEnvironment; import com.google.devtools.build.lib.query2.PostAnalysisQueryEnvironment.TopLevelConfigurations; +import com.google.devtools.build.lib.query2.common.CommonQueryOptions; import com.google.devtools.build.lib.query2.engine.QueryEvalResult; import com.google.devtools.build.lib.query2.engine.QueryException; import com.google.devtools.build.lib.query2.engine.QueryExpression; @@ -90,7 +91,7 @@ public void process( } try (QueryRuntimeHelper queryRuntimeHelper = - env.getRuntime().getQueryRuntimeHelperFactory().create(env)) { + env.getRuntime().getQueryRuntimeHelperFactory().create(env, getQueryOptions(env))) { doPostAnalysisQuery( request, env, @@ -131,6 +132,8 @@ public void process( } } + protected abstract CommonQueryOptions getQueryOptions(CommandEnvironment env); + protected abstract PostAnalysisQueryEnvironment getQueryEnvironment( BuildRequest request, CommandEnvironment env, diff --git a/src/main/java/com/google/devtools/build/lib/query2/common/CommonQueryOptions.java b/src/main/java/com/google/devtools/build/lib/query2/common/CommonQueryOptions.java index 307191ae09e1d9..efc861d797ccaa 100644 --- a/src/main/java/com/google/devtools/build/lib/query2/common/CommonQueryOptions.java +++ b/src/main/java/com/google/devtools/build/lib/query2/common/CommonQueryOptions.java @@ -373,6 +373,10 @@ public AspectResolutionModeConverter() { + "applicable to --output=graph.") public boolean graphFactored; + /////////////////////////////////////////////////////////// + // INPUT / OUTPUT OPTIONS // + /////////////////////////////////////////////////////////// + @Option( name = "query_file", defaultValue = "", @@ -382,4 +386,15 @@ public AspectResolutionModeConverter() { "If set, query will read the query from the file named here, rather than on the command " + "line. It is an error to specify a file here as well as a command-line query.") public String queryFile; + + @Option( + name = "output_file", + defaultValue = "", + documentationCategory = OptionDocumentationCategory.QUERY, + effectTags = {OptionEffectTag.TERMINAL_OUTPUT}, + help = + "When specified, query results will be written directly to this file, and nothing will be" + + " printed to Bazel's standard output stream (stdout). In benchmarks, this is" + + " generally faster than bazel query > file.") + public String outputFile; } diff --git a/src/main/java/com/google/devtools/build/lib/runtime/QueryRuntimeHelper.java b/src/main/java/com/google/devtools/build/lib/runtime/QueryRuntimeHelper.java index e7bac1a87c2e29..1baea301e5c47f 100644 --- a/src/main/java/com/google/devtools/build/lib/runtime/QueryRuntimeHelper.java +++ b/src/main/java/com/google/devtools/build/lib/runtime/QueryRuntimeHelper.java @@ -15,10 +15,14 @@ import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; +import com.google.common.base.Strings; +import com.google.devtools.build.lib.query2.common.CommonQueryOptions; import com.google.devtools.build.lib.server.FailureDetails; import com.google.devtools.build.lib.server.FailureDetails.FailureDetail; import com.google.devtools.build.lib.server.FailureDetails.Query; import com.google.devtools.build.lib.server.FailureDetails.Query.Code; +import com.google.devtools.build.lib.vfs.Path; +import java.io.IOException; import java.io.OutputStream; /** @@ -47,7 +51,8 @@ public interface QueryRuntimeHelper extends AutoCloseable { /** Factory for {@link QueryRuntimeHelper} instances. */ interface Factory { - QueryRuntimeHelper create(CommandEnvironment env) throws QueryRuntimeHelperException; + QueryRuntimeHelper create(CommandEnvironment env, CommonQueryOptions options) + throws QueryRuntimeHelperException; } /** @@ -55,6 +60,8 @@ interface Factory { * {@link CommandEnvironment} instance's stdout. * *

This is intended to be the default {@link Factory}. + * + *

If {@code --output_file} is set, the stdout is redirected to the defined path value instead. */ class StdoutQueryRuntimeHelperFactory implements Factory { public static final StdoutQueryRuntimeHelperFactory INSTANCE = @@ -63,8 +70,14 @@ class StdoutQueryRuntimeHelperFactory implements Factory { private StdoutQueryRuntimeHelperFactory() {} @Override - public QueryRuntimeHelper create(CommandEnvironment env) { - return createInternal(env.getReporter().getOutErr().getOutputStream()); + public QueryRuntimeHelper create(CommandEnvironment env, CommonQueryOptions options) + throws QueryRuntimeHelperException { + if (Strings.isNullOrEmpty(options.outputFile)) { + return createInternal(env.getReporter().getOutErr().getOutputStream()); + } else { + return FileQueryRuntimeHelper.create( + env.getWorkingDirectory().getRelative(options.outputFile)); + } } public QueryRuntimeHelper createInternal(OutputStream stdoutOutputStream) { @@ -91,6 +104,49 @@ public void afterQueryOutputIsWritten() {} @Override public void close() {} } + + /** + * A {@link QueryRuntimeHelper} that wraps a {@link java.io.FileOutputStream} instead of writing + * to standard out, for improved performance. + */ + public static class FileQueryRuntimeHelper implements QueryRuntimeHelper { + private final Path path; + private final OutputStream out; + + private FileQueryRuntimeHelper(Path path) throws IOException { + this.path = path; + this.out = path.getOutputStream(); + } + + public static FileQueryRuntimeHelper create(Path path) throws QueryRuntimeHelperException { + try { + return new FileQueryRuntimeHelper(path); + } catch (IOException e) { + throw new QueryRuntimeHelperException( + "Could not open query output file " + path.getPathString(), + Code.QUERY_OUTPUT_WRITE_FAILURE, + e); + } + } + + @Override + public OutputStream getOutputStreamForQueryOutput() { + return out; + } + + @Override + public void afterQueryOutputIsWritten() {} + + @Override + public void close() throws QueryRuntimeHelperException { + try { + out.close(); + } catch (IOException e) { + throw new QueryRuntimeHelperException( + "Could not close query output file " + path, Code.QUERY_OUTPUT_WRITE_FAILURE, e); + } + } + } } /** Describes what went wrong in {@link QueryRuntimeHelper}. */ diff --git a/src/main/java/com/google/devtools/build/lib/runtime/commands/QueryEnvironmentBasedCommand.java b/src/main/java/com/google/devtools/build/lib/runtime/commands/QueryEnvironmentBasedCommand.java index 8b89f9f61fc078..3a9b5500085898 100644 --- a/src/main/java/com/google/devtools/build/lib/runtime/commands/QueryEnvironmentBasedCommand.java +++ b/src/main/java/com/google/devtools/build/lib/runtime/commands/QueryEnvironmentBasedCommand.java @@ -167,7 +167,7 @@ private BlazeCommandResult execInternal(CommandEnvironment env, OptionsParsingRe .getLabelPrinter(starlarkSemantics, mainRepoTargetParser.getRepoMapping()); try (QueryRuntimeHelper queryRuntimeHelper = - env.getRuntime().getQueryRuntimeHelperFactory().create(env)) { + env.getRuntime().getQueryRuntimeHelperFactory().create(env, queryOptions)) { Either result; try (AbstractBlazeQueryEnvironment queryEnv = newQueryEnvironment( diff --git a/src/test/shell/integration/bazel_query_test.sh b/src/test/shell/integration/bazel_query_test.sh index b7a428dce864fb..fb8c9bfc6b8c58 100755 --- a/src/test/shell/integration/bazel_query_test.sh +++ b/src/test/shell/integration/bazel_query_test.sh @@ -42,7 +42,7 @@ fi source "$(rlocation "io_bazel/src/test/shell/integration_test_setup.sh")" \ || { echo "integration_test_setup.sh not found!" >&2; exit 1; } -# `uname` returns the current platform, e.g "MSYS_NT-10.0" or "Linux". +# `uname` returns the current platform, e.g. "MSYS_NT-10.0" or "Linux". # `tr` converts all upper case letters to lower case. # `case` matches the result if the `uname | tr` expression to string prefixes # that use the same wildcards as names do in Bash, i.e. "msys*" matches strings @@ -82,6 +82,22 @@ EOF expect_log "//peach:harken" } +function test_output_to_file() { + rm -rf peach + mkdir -p peach + cat > peach/BUILD < $TEST_TMPDIR/query_stdout + + expect_log "//peach:brighton" + expect_log "//peach:harken" + + assert_equals "" "$(<$TEST_TMPDIR/query_stdout)" +} + function test_invalid_query_fails_parsing() { bazel query 'deps("--bad_target_name_from_bad_script")' >& "$TEST_log" \ && fail "Expected failure"