Skip to content

Speed up streamed-proto query output by distributing work to multiple threads #12

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 10 commits into
base: master
Choose a base branch
from

Conversation

arvi18
Copy link

@arvi18 arvi18 commented Apr 27, 2025

This is a proposed fix for bazelbuild#24304

This speeds up a fully warm bazel query ... by 54%, reducing wall time from 1m49s to 50s

Current state:

$ time bazel query '...' --output=streamed_proto > queryoutput4.streamedproto

real    1m48.768s
user    0m27.410s
sys     0m19.646s

This PR:

$ time bazel query '...' --output=streamed_proto > queryoutput5.streamedproto

real    0m49.938s
user    0m22.897s
sys     0m16.161s

💁‍♂️ Note: when combined with bazelbuild#24298, total wall time is 37s, an overall reduction of 66%.

Summary by CodeRabbit

  • New Features
    • Improved performance for streaming large sets of proto targets by introducing asynchronous, chunked processing and streaming.
  • Bug Fixes
    • Enhanced reliability and responsiveness when handling large queries by managing concurrency and backpressure during output streaming.

@arvi18
Copy link
Author

arvi18 commented Apr 27, 2025

Do you think this is something we could add behind a flag to start off? I agree it's useful to have, but it also comes with misc subtleties that are hard to reason about all at once, so having an escape hatch in case we get it wrong the first time would be nice.

@arvi18
Copy link
Author

arvi18 commented Apr 27, 2025

I had actually been thinking of more of an opt-in that possibly makes it more clear it's experimental, depending on how solid we can get it on the first try, then eventually flip the default. The whole parallel-streaming-query thing gets tricky when we need to get exception handling involved, an opt-in will let us get code in and iterate with less risk of breakages.

Copy link

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hello @arvi18, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

Summary of Changes

This pull request addresses issue bazelbuild#24304 by speeding up the bazel query ... command when used with the --output=streamed_proto option. The changes introduce multi-threading to distribute the work of writing the proto output, resulting in a reported 54% reduction in wall time for a fully warm query. The original implementation was single-threaded, and this change leverages a ForkJoinPool to parallelize the writing of target proto buffers to the output stream.

Highlights

  • Performance Improvement: The primary goal of this PR is to improve the performance of bazel query with --output=streamed_proto by distributing the workload across multiple threads.
  • Multi-threading Implementation: A ForkJoinPool is used to parallelize the process of converting Target objects to proto buffers and writing them to the output stream. The number of threads defaults to the number of available processors.
  • Chunking and Queuing: The input Iterable<Target> is partitioned into chunks of 500 targets each. These chunks are then submitted to the ForkJoinPool via a LinkedBlockingQueue to be processed in parallel.
  • Error Handling: WrappedIOException and WrappedInterruptedException are introduced to handle exceptions thrown during the parallel processing and re-thrown as RuntimeExceptions.

Changelog

  • src/main/java/com/google/devtools/build/lib/query2/query/output/StreamedProtoOutputFormatter.java
    • Introduces multi-threading using ForkJoinPool to parallelize proto output writing.
    • Partitions the input targets into chunks for parallel processing.
    • Uses a LinkedBlockingQueue to manage the chunks to be processed by the thread pool.
    • Adds WrappedIOException and WrappedInterruptedException to handle exceptions during parallel processing.
    • Implements writeTargetsDelimitedToByteArrays to convert a list of targets to byte arrays in parallel.
    • Implements toProto to convert a target to a proto object.
    • Implements writeDelimited to write a proto object to a byte array in a delimited format.
    • Adds constants for MAX_CHUNKS_IN_QUEUE and TARGETS_PER_CHUNK to control the chunking and queuing behavior.
Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in issue comments and review comments.
Help /gemini help Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.


A query's slow pace,
Threads now race, time and space,
Swift proto takes flight.

Footnotes

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution.

Copy link

coderabbitai bot commented Apr 27, 2025

Walkthrough

The StreamedProtoOutputFormatter class was refactored to change how it serializes and writes proto targets to an output stream. Instead of synchronously writing each target, the new implementation partitions targets into chunks, serializes each chunk asynchronously using a ForkJoinPool, and manages these operations with a bounded blocking queue. The serialized results are then written to the output stream in order by a dedicated consumer. Several helper methods and exception wrapper classes were introduced to support this concurrent, chunked, and streamed processing model.

Changes

Cohort / File(s) Change Summary
Core Formatter Refactor
src/main/java/com/google/devtools/build/lib/query2/query/output/StreamedProtoOutputFormatter.java
Refactored processOutput to use chunked, concurrent serialization and streaming of proto targets; added private helper methods for chunk serialization and proto conversion; rewrote length-delimited proto serialization method; introduced private static runtime exception wrappers for IOException and InterruptedException; updated internal logic of createPostFactoStreamCallback to implement producer-consumer model with ForkJoinPool and bounded queue.

Sequence Diagram(s)

sequenceDiagram
    participant Client
    participant OutputFormatterCallback
    participant ForkJoinPool
    participant BlockingQueue
    participant OutputStream

    Client->>OutputFormatterCallback: processOutput(targets)
    OutputFormatterCallback->>ForkJoinPool: Submit chunked serialization tasks
    ForkJoinPool-->>BlockingQueue: Place futures for each chunk
    loop For each chunk future in order
        BlockingQueue->>OutputFormatterCallback: Take next completed future
        OutputFormatterCallback->>OutputStream: Write serialized bytes
    end
    OutputFormatterCallback-->>Client: Completion
Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~15 minutes

Poem

In a warren of code, the bytes now stream,
Chunks hop along, as in a dream.
Forks and joins, a queue in tow,
Proto targets in a rabbit row.
Streams are swift, exceptions wrapped tight—
The formatter’s future is running bright!
🐇✨

Note

⚡️ Unit Test Generation is now available in beta!

Learn more here, or try it out under "Finishing Touches" below.

✨ Finishing Touches
  • 📝 Generate Docstrings
🧪 Generate unit tests
  • Create PR with unit tests
  • Post copyable unit tests in a comment
  • Commit unit tests in branch clone-parallel-streamed-proto

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share
🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

  • Review comments: Directly reply to a review comment made by CodeRabbit. Example:
    • I pushed a fix in commit <commit_id>, please review it.
    • Explain this complex logic.
    • Open a follow-up GitHub issue for this discussion.
  • Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
    • @coderabbitai explain this code block.
    • @coderabbitai modularize this function.
  • PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
    • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
    • @coderabbitai read src/utils.ts and explain its main purpose.
    • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
    • @coderabbitai help me debug CodeRabbit configuration file.

Support

Need help? Create a ticket on our support page for assistance with any issues or questions.

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

CodeRabbit Commands (Invoked using PR comments)

  • @coderabbitai pause to pause the reviews on a PR.
  • @coderabbitai resume to resume the paused reviews.
  • @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
  • @coderabbitai full review to do a full review from scratch and review all the files again.
  • @coderabbitai summary to regenerate the summary of the PR.
  • @coderabbitai generate docstrings to generate docstrings for this PR.
  • @coderabbitai generate sequence diagram to generate a sequence diagram of the changes in this PR.
  • @coderabbitai generate unit tests to generate unit tests for this PR.
  • @coderabbitai resolve resolve all the CodeRabbit review comments.
  • @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
  • @coderabbitai help to get help.

Other keywords and placeholders

  • Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
  • Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
  • Add @coderabbitai anywhere in the PR title to generate the title automatically.

CodeRabbit Configuration File (.coderabbit.yaml)

  • You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
  • Please see the configuration documentation for more information.
  • If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json

Documentation and Community

  • Visit our Documentation for detailed information on how to use CodeRabbit.
  • Join our Discord Community to get help, request features, and share feedback.
  • Follow us on X/Twitter for updates and announcements.

Copy link

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces multi-threading to the streamed proto output formatter, significantly improving query performance. The approach of using a ForkJoinPool and a blocking queue seems reasonable for distributing the work. However, there are a few areas that could be improved, particularly around error handling and resource management.

Summary of Findings

  • Error Propagation from Executor: The code uses RuntimeException to wrap exceptions from the executor, losing valuable information about the original exception. It's crucial to propagate the original exception to provide better debugging insights.
  • Resource Management of ForkJoinPool: While the ForkJoinPool is created within a try-with-resources block, ensuring proper shutdown, the tasks submitted to the pool might still be running after the block completes if they are long-running or blocked. Consider explicitly shutting down the executor and awaiting termination to ensure all tasks are completed before proceeding.
  • Wrapped Exception Classes: The WrappedIOException and WrappedInterruptedException classes are simple wrappers around checked exceptions. Consider whether these wrappers are truly necessary, or if the exceptions can be handled directly.

Merge Readiness

The performance improvements introduced by this pull request are significant, making it a valuable contribution. However, the error handling and resource management aspects need to be addressed before merging. I recommend resolving the issues raised in the review comments, particularly regarding exception propagation and ForkJoinPool shutdown. I am unable to approve this pull request, and other reviewers should review and approve this code before merging.

Comment on lines +79 to +81
// TODO: figure out what might be in here and propagate
throw new RuntimeException(e);
}

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

This TODO indicates that the exception handling here is incomplete. It's important to propagate the original ExecutionException's cause to provide more context about the failure. Consider using e.getCause() to get the underlying exception and wrap it in a more specific exception type, or rethrow it if appropriate.

Also, consider logging the exception with sufficient context to aid debugging.

                      // TODO: figure out what might be in here and propagate
                      Throwable cause = e.getCause();
                      if (cause instanceof IOException) {
                        throw new WrappedIOException((IOException) cause);
                      } else if (cause instanceof InterruptedException) {
                        throw new WrappedInterruptedException((InterruptedException) cause);
                      } else {
                        throw new RuntimeException("Error during target processing", cause);
                      }

Comment on lines +94 to +96
// TODO: propagate
throw new RuntimeException(e);
}

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Similar to the previous TODO, this one also requires proper exception propagation. The ExecutionException should be unwrapped to get the root cause and either rethrown as a more specific exception or wrapped in a RuntimeException with the cause attached. Logging the exception is also recommended.

        } catch (ExecutionException e) {
          // TODO: propagate
          Throwable cause = e.getCause();
          if (cause instanceof IOException) {
            throw new WrappedIOException((IOException) cause);
          } else if (cause instanceof InterruptedException) {
            throw new WrappedInterruptedException((InterruptedException) cause);
          } else {
            throw new RuntimeException("Error during target processing", cause);
          }

Comment on lines +88 to +90
stillAddingTargetsToQueue.set(false);
}
}

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

After setting stillAddingTargetsToQueue to false, it's good practice to explicitly shutdown the executor and wait for its termination. This ensures that all submitted tasks are completed before proceeding. Otherwise, there is a risk that the executor will not shutdown cleanly, and resources may be leaked. Consider adding executor.shutdown() and executor.awaitTermination() calls within the finally block.

          } finally {
            stillAddingTargetsToQueue.set(false);
            executor.shutdown(); // Initiate shutdown
            try {
              if (!executor.awaitTermination(60, TimeUnit.SECONDS)) { // Wait for tasks to complete
                System.err.println("ForkJoinPool did not terminate in the specified timeout.");
                // Optionally, log the state of the tasks that are still running.
              }
            } catch (InterruptedException ie) {
              // (Re-)Cancel if current thread also interrupted
              executor.shutdownNow();
              // Preserve interrupt status
              Thread.currentThread().interrupt();
            }
          }

Comment on lines +132 to +134
@Override
public IOException getCause() {
return (IOException) super.getCause();

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Consider whether wrapping the IOException in a RuntimeException is truly necessary. If the calling code can handle IOException directly, it might be better to avoid the wrapping. If wrapping is necessary, ensure that the original exception is accessible via getCause().

    @Override
    public IOException getCause() {
      return cause;
    }

Comment on lines +144 to +146
public InterruptedException getCause() {
return (InterruptedException) super.getCause();
}

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Similar to WrappedIOException, evaluate the necessity of wrapping InterruptedException. If the calling code can handle it directly, avoid wrapping. If wrapping is necessary, ensure the original exception is accessible via getCause().

    @Override
    public InterruptedException getCause() {
      return cause;
    }

Copy link

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 2

🧹 Nitpick comments (1)
src/main/java/com/google/devtools/build/lib/query2/query/output/StreamedProtoOutputFormatter.java (1)

92-96: Propagate checked exceptions instead of wrapping them in generic RuntimeException

writeAllTargetsFuture.get() can surface WrappedIOException / WrappedInterruptedException. By catching ExecutionException and blindly re-wrapping, callers lose the ability to handle IOException or honour the method’s declared InterruptedException.

Suggestion:

-        } catch (ExecutionException e) {
-          // TODO: propagate
-          throw new RuntimeException(e);
+        } catch (ExecutionException e) {
+          Throwable cause = e.getCause();
+          if (cause instanceof WrappedIOException ioe) {
+            throw ioe.getCause();          // preserves declared throws IOException
+          }
+          if (cause instanceof WrappedInterruptedException ie) {
+            Thread.currentThread().interrupt();
+            throw ie.getCause();
+          }
+          throw new RuntimeException(cause);
         }

This keeps the public contract intact and avoids masking interruption.

📜 Review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between c0c98a2 and 5fc8b13.

📒 Files selected for processing (1)
  • src/main/java/com/google/devtools/build/lib/query2/query/output/StreamedProtoOutputFormatter.java (2 hunks)
🔇 Additional comments (1)
src/main/java/com/google/devtools/build/lib/query2/query/output/StreamedProtoOutputFormatter.java (1)

63-90: Potential producer/consumer termination race

stillAddingTargetsToQueue is set to false only in the finally block after the producer loop. If partialResult is empty, the queue remains empty and the consumer immediately blocks on take() forever, preventing writeAllTargetsFuture from completing.

Consider signalling completion explicitly:

// after producing chunks
targetQueue.put(POISON_PILL_FUTURE); // Future that returns an empty list
...
while (true) {
  Future<List<byte[]>> f = targetQueue.take();
  if (f == POISON_PILL_FUTURE) break;
  ...
}

or a BlockingQueue with null sentinel, ensuring graceful shutdown even for empty inputs.

Comment on lines +55 to +72
new ForkJoinPool(
Runtime.getRuntime().availableProcessors(),
ForkJoinPool.defaultForkJoinWorkerThreadFactory,
null,
// we use asyncMode to ensure the queue is processed FIFO, which maximizes
// throughput
true)) {
var targetQueue = new LinkedBlockingQueue<Future<List<byte[]>>>(MAX_CHUNKS_IN_QUEUE);
var stillAddingTargetsToQueue = new AtomicBoolean(true);
writeAllTargetsFuture =
executor.submit(
() -> {
try {
while (stillAddingTargetsToQueue.get() || !targetQueue.isEmpty()) {
Future<List<byte[]>> targets = targetQueue.take();
for (byte[] target : targets.get()) {
out.write(target);
}
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

ForkJoinPool worker is blocking on take()/get() – high risk of thread starvation

The consumer task runs inside the same ForkJoinPool that processes the producer tasks yet:

  • performs unmanaged blocking (targetQueue.take() and Future.get()),
  • uses a bounded queue that may fill up, causing the producer thread (caller of processOutput) to block on put().

Because ForkJoinPool counts blocked workers toward parallelism, the pool can dead-lock or under-utilise CPUs when many long-running producer tasks occupy the limited workers while the single consumer waits, or vice-versa.

Recommended approaches:

  1. Start the consumer on a dedicated thread (e.g. Executors.newSingleThreadExecutor) outside the FJP, or
  2. Wrap blocking calls with ForkJoinPool.managedBlock, or
  3. Replace the FJP entirely with a plain ExecutorService that tolerates blocking.

This will eliminate the starvation risk and make behaviour more predictable.

@visz11
Copy link
Collaborator

visz11 commented Jul 26, 2025

/refacto-test

Copy link

refacto-test bot commented Jul 26, 2025

Refacto is reviewing this PR. Please wait for the review comments to be posted.

throw new WrappedInterruptedException(e);
} catch (IOException e) {
throw new WrappedIOException(e);
} catch (ExecutionException e) {
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thread Safety Issue: The ExecutionException handling is inadequate, using a generic RuntimeException with a TODO comment. This can mask important exceptions during task execution, potentially leading to resource leaks or thread pool exhaustion.

Consider unwrapping and properly handling the cause:

catch (ExecutionException e) {
  Throwable cause = e.getCause();
  if (cause instanceof IOException) {
    throw new WrappedIOException((IOException) cause);
  } else if (cause instanceof InterruptedException) {
    Thread.currentThread().interrupt(); // Preserve interrupt status
    throw new WrappedInterruptedException((InterruptedException) cause);
  } else if (cause instanceof RuntimeException) {
    throw (RuntimeException) cause;
  } else {
    throw new RuntimeException("Error processing targets", cause);
  }
}

@Override
public void processOutput(Iterable<Target> partialResult)
throws IOException, InterruptedException {
for (Target target : partialResult) {
toTargetProtoBuffer(target, labelPrinter).writeDelimitedTo(out);
ForkJoinTask<?> writeAllTargetsFuture;
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Resource Management Issue: The ForkJoinPool is created with parallelism equal to all available processors without any safeguards against resource exhaustion. On many-core systems, this could lead to excessive thread creation and CPU contention.

Consider limiting the maximum parallelism and adding an exception handler:

// Limit the number of threads to avoid resource exhaustion
int parallelism = Math.min(Runtime.getRuntime().availableProcessors(), 4);
try (ForkJoinPool executor = new ForkJoinPool(
    parallelism,
    ForkJoinPool.defaultForkJoinWorkerThreadFactory,
    (thread, throwable) -> {
        // Log uncaught exceptions in worker threads
        System.err.println("Worker thread exception: " + throwable);
    },
    // we use asyncMode to ensure the queue is processed FIFO
    true)) {

}
});
try {
for (List<Target> targets : Iterables.partition(partialResult, TARGETS_PER_CHUNK)) {
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Concurrency Issue: The blocking targetQueue.put() operation could lead to thread starvation if the queue becomes full. There's no timeout or circuit-breaking mechanism to handle slow or stuck tasks.

Consider using a non-blocking approach with timeouts:

// Use offer with timeout to prevent indefinite blocking
Future<List<byte[]>> task = executor.submit(() -> writeTargetsDelimitedToByteArrays(targets));
if (!targetQueue.offer(task, 5, TimeUnit.SECONDS)) {
  // If queue is full for too long, process this batch directly to make progress
  for (byte[] target : task.get(30, TimeUnit.SECONDS)) {
    out.write(target);
  }
}

…t/StreamedProtoOutputFormatter.java

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
@visz11
Copy link
Collaborator

visz11 commented Jul 29, 2025

/refacto-test

Copy link

refacto-test bot commented Jul 29, 2025

Refacto is reviewing this PR. Please wait for the review comments to be posted.

Copy link

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 0

♻️ Duplicate comments (5)
src/main/java/com/google/devtools/build/lib/query2/query/output/StreamedProtoOutputFormatter.java (5)

53-61: Address resource management and parallelism concerns from previous reviews

The ForkJoinPool creation has several issues that were previously identified but remain unaddressed:

  1. No upper bound on parallelism (could exhaust resources on many-core systems)
  2. No exception handler for uncaught worker thread exceptions

Apply this fix to address the resource management issues:

-try (ForkJoinPool executor =
-    new ForkJoinPool(
-        Runtime.getRuntime().availableProcessors(),
-        ForkJoinPool.defaultForkJoinWorkerThreadFactory,
-        null,
-        // we use asyncMode to ensure the queue is processed FIFO, which maximizes
-        // throughput
-        true)) {
+int parallelism = Math.min(Runtime.getRuntime().availableProcessors(), 8);
+try (ForkJoinPool executor =
+    new ForkJoinPool(
+        parallelism,
+        ForkJoinPool.defaultForkJoinWorkerThreadFactory,
+        (thread, throwable) -> {
+            System.err.println("Worker thread exception in query processing: " + throwable);
+        },
+        // we use asyncMode to ensure the queue is processed FIFO, which maximizes
+        // throughput
+        true)) {

64-82: Critical threading architecture issue - consumer runs in same ForkJoinPool

The consumer task is submitted to the same ForkJoinPool that processes the producer tasks, creating a high risk of deadlock and thread starvation as previously identified. The consumer performs blocking operations (take(), get()) which can starve the work-stealing threads.

Move the consumer to a dedicated thread outside the ForkJoinPool:

+ExecutorService consumerExecutor = Executors.newSingleThreadExecutor();
+try {
  var targetQueue = new LinkedBlockingQueue<Future<List<byte[]>>>(MAX_CHUNKS_IN_QUEUE);
  var stillAddingTargetsToQueue = new AtomicBoolean(true);
- writeAllTargetsFuture =
-     executor.submit(
+ writeAllTargetsFuture =
+     consumerExecutor.submit(
          () -> {
            // consumer logic remains the same
          });
+} finally {
+  consumerExecutor.shutdown();
+  if (!consumerExecutor.awaitTermination(60, TimeUnit.SECONDS)) {
+    consumerExecutor.shutdownNow();
+  }
+}

78-81: Incomplete exception handling in consumer task

The TODO comment indicates incomplete exception propagation. This masks important failures and can lead to silent data loss or resource leaks.

Properly unwrap and propagate the ExecutionException cause:

-} catch (ExecutionException e) {
-  // TODO: figure out what might be in here and propagate
-  throw new RuntimeException(e);
+} catch (ExecutionException e) {
+  Throwable cause = e.getCause();
+  if (cause instanceof IOException) {
+    throw new WrappedIOException((IOException) cause);
+  } else if (cause instanceof InterruptedException) {
+    Thread.currentThread().interrupt();
+    throw new WrappedInterruptedException((InterruptedException) cause);
+  } else if (cause instanceof RuntimeException) {
+    throw (RuntimeException) cause;
+  } else {
+    throw new RuntimeException("Error processing targets", cause);
+  }

84-89: Potential blocking issue with bounded queue

The targetQueue.put() operation can block indefinitely if the consumer falls behind, as previously identified. This could lead to thread starvation.

Use a timeout-based approach to prevent indefinite blocking:

-for (List<Target> targets : Iterables.partition(partialResult, TARGETS_PER_CHUNK)) {
-  targetQueue.put(executor.submit(() -> writeTargetsDelimitedToByteArrays(targets)));
-}
+for (List<Target> targets : Iterables.partition(partialResult, TARGETS_PER_CHUNK)) {
+  Future<List<byte[]>> task = executor.submit(() -> writeTargetsDelimitedToByteArrays(targets));
+  if (!targetQueue.offer(task, 30, TimeUnit.SECONDS)) {
+    throw new RuntimeException("Queue timeout - consumer may be stuck");
+  }
+}

94-96: Duplicate incomplete exception handling

Same issue as in the consumer task - ExecutionException is not properly unwrapped and propagated.

Apply the same fix as for the consumer task:

-} catch (ExecutionException e) {
-  // TODO: propagate
-  throw new RuntimeException(e);
+} catch (ExecutionException e) {
+  Throwable cause = e.getCause();
+  if (cause instanceof IOException) {
+    throw (IOException) cause;
+  } else if (cause instanceof InterruptedException) {
+    Thread.currentThread().interrupt();
+    throw (InterruptedException) cause;
+  } else if (cause instanceof RuntimeException) {
+    throw (RuntimeException) cause;
+  } else {
+    throw new RuntimeException("Error waiting for consumer completion", cause);
+  }
🧹 Nitpick comments (1)
src/main/java/com/google/devtools/build/lib/query2/query/output/StreamedProtoOutputFormatter.java (1)

45-47: Consider making chunk size and queue capacity configurable

The hardcoded constants may not be optimal for all environments:

  • TARGETS_PER_CHUNK = 500 might be too small/large depending on target complexity
  • MAX_CHUNKS_IN_QUEUE = processors * 2 could cause memory issues with large targets

Consider making these configurable via system properties or constructor parameters:

-private static final int MAX_CHUNKS_IN_QUEUE = Runtime.getRuntime().availableProcessors() * 2;
-private static final int TARGETS_PER_CHUNK = 500;
+private static final int MAX_CHUNKS_IN_QUEUE = Integer.getInteger("bazel.query.max_chunks", Runtime.getRuntime().availableProcessors() * 2);
+private static final int TARGETS_PER_CHUNK = Integer.getInteger("bazel.query.targets_per_chunk", 500);
📜 Review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 5fc8b13 and ec1bf6e.

📒 Files selected for processing (1)
  • src/main/java/com/google/devtools/build/lib/query2/query/output/StreamedProtoOutputFormatter.java (2 hunks)
🔇 Additional comments (6)
src/main/java/com/google/devtools/build/lib/query2/query/output/StreamedProtoOutputFormatter.java (6)

16-21: LGTM - Import additions support the concurrent implementation

The added imports (Iterables, Build, CodedOutputStream, List, concurrent utilities, and AtomicBoolean) are all necessary for the new multi-threaded implementation.

Also applies to: 25-27


99-101: LGTM - Clean helper method implementation

The writeTargetsDelimitedToByteArrays method correctly maps targets to delimited byte arrays using the fixed writeDelimited method.


103-109: LGTM - Proper exception wrapping for concurrent context

The toProto method correctly wraps InterruptedException in a runtime exception, which is necessary for use within the ForkJoinPool's lambda expressions.


113-138: Excellent fix for length-delimited proto serialization

The previous implementation had a critical bug where the length prefix was never actually written, making the output unparseable. The new implementation correctly:

  1. Writes the varint length prefix using a separate CodedOutputStream
  2. Writes the message bytes after the prefix
  3. Properly flushes both streams

This addresses the critical issue identified in previous reviews where mergeDelimitedFrom would fail.


140-149: LGTM - Exception wrapper provides proper cause access

The WrappedIOException class correctly extends RuntimeException and overrides getCause() to return the specific IOException type, addressing the concern from previous reviews.


151-160: LGTM - Consistent exception wrapper implementation

The WrappedInterruptedException class follows the same pattern as WrappedIOException, providing proper access to the underlying InterruptedException.

@visz11
Copy link
Collaborator

visz11 commented Jul 31, 2025

/refacto-test

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants