Skip to content

Commit

Permalink
Add bazel query --output_file option, which writes query results di…
Browse files Browse the repository at this point in the history
…rectly to a file

This is a proposed fix for #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 #24305, total wall time is 37s, an overall reduction of 66%._

Closes #24298.

PiperOrigin-RevId: 700583890
Change-Id: Ic13f0611aca60c2ce8641e72a0fcfc330f13c803
  • Loading branch information
keithl-stripe authored and copybara-github committed Nov 27, 2024
1 parent c79749a commit 791e1f7
Show file tree
Hide file tree
Showing 9 changed files with 111 additions and 9 deletions.
2 changes: 1 addition & 1 deletion scripts/bash_completion_test.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions src/main/java/com/google/devtools/build/lib/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -403,6 +403,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",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -131,6 +132,8 @@ public void process(
}
}

protected abstract CommonQueryOptions getQueryOptions(CommandEnvironment env);

protected abstract PostAnalysisQueryEnvironment<T> getQueryEnvironment(
BuildRequest request,
CommandEnvironment env,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -373,6 +373,10 @@ public AspectResolutionModeConverter() {
+ "applicable to --output=graph.")
public boolean graphFactored;

///////////////////////////////////////////////////////////
// INPUT / OUTPUT OPTIONS //
///////////////////////////////////////////////////////////

@Option(
name = "query_file",
defaultValue = "",
Expand All @@ -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 <code>bazel query &gt; file</code>.")
public String outputFile;
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;

/**
Expand Down Expand Up @@ -47,14 +51,17 @@ 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;
}

/**
* A {@link Factory} for {@link StdoutQueryRuntimeHelper} instances that simply wrap the given
* {@link CommandEnvironment} instance's stdout.
*
* <p>This is intended to be the default {@link Factory}.
*
* <p>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 =
Expand All @@ -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) {
Expand All @@ -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}. */
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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<BlazeCommandResult, QueryEvalResult> result;
try (AbstractBlazeQueryEnvironment<Target> queryEnv =
newQueryEnvironment(
Expand Down
18 changes: 17 additions & 1 deletion src/test/shell/integration/bazel_query_test.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -82,6 +82,22 @@ EOF
expect_log "//peach:harken"
}

function test_output_to_file() {
rm -rf peach
mkdir -p peach
cat > peach/BUILD <<EOF
sh_library(name='brighton', deps=[':harken'])
sh_library(name='harken')
EOF

bazel query 'deps(//peach:brighton)' --output_file=$TEST_log > $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"
Expand Down

0 comments on commit 791e1f7

Please sign in to comment.