Optimize Kafka IO: Remove Guava Stopwatch overhead from consumer.poll() loop

[junaiddshaukat]

Fixes #37704

What happened?

This PR optimizes the Kafka IO readers (ReadFromKafkaDoFn and KafkaUnboundedReader) by removing the Guava Stopwatch and the custom RPC latency metric tracking from the hot consumer.poll() loop.

The Problem:
Guava's Stopwatch relies on System.nanoTime(), which is a relatively expensive system call. In the Kafka IO module, this timer was being started and stopped repeatedly inside the consumer.poll() loop. When Kafka prefetches records, this loop executes extremely rapidly, causing the Stopwatch to introduce significant and unnecessary CPU overhead.

The Solution:

1. Removed Stopwatch: Replaced the Stopwatch timeout tracking with a lightweight System.currentTimeMillis() calculation to track the remainingTimeout.
2. Removed Custom RPC Metrics: Removed the updateSuccessfulRpcMetrics calls and associated histograms (KafkaMetrics, KafkaSinkMetrics).
3. Why this is safe: The custom Beam RPC latency metric is redundant. Kafka already natively tracks and exposes this data via JMX metrics (specifically fetch-latency-avg), so we can safely drop the custom implementation to regain performance.

Changes made:

• Updated ReadFromKafkaDoFn.java and KafkaUnboundedReader.java to use System.currentTimeMillis() for timeout tracking instead of Stopwatch.
• Cleaned up dead code in KafkaMetrics.java and KafkaSinkMetrics.java by removing RPC_LATENCY histograms and updateSuccessfulRpcMetrics.
• Updated KafkaMetricsTest.java and KafkaSinkMetricsTest.java to reflect the removed metrics.

Testing Done

• Compiled successfully.
• Ran and passed the full Kafka IO test suite (./gradlew :sdks:java:io:kafka:test).
• Verified that timeout logic in the poll loop still functions correctly using java.time.Duration and System.currentTimeMillis().

---

Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:

• Mention the appropriate issue in your description (for example: addresses #123), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, comment fixes #<ISSUE NUMBER> instead.
• Update CHANGES.md with noteworthy changes.
• If this contribution is large, please file an Apache Individual Contributor License Agreement.

See the Contributor Guide for more tips on how to make review process smoother.

To check the build health, please visit https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md

GitHub Actions Tests Status (on master branch)

See CI.md for more information about GitHub Actions CI or the workflows README to see a list of phrases to trigger workflows.

[gemini-code-assist]

Summary of Changes

Hello @junaiddshaukat, I'm Gemini Code Assist¹! 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!

This pull request significantly enhances the performance of Kafka IO readers by addressing and removing two sources of unnecessary CPU overhead. By replacing an expensive timer mechanism with a more efficient alternative and eliminating redundant custom RPC latency metrics, the Kafka consumer's poll() loop now executes with reduced latency, especially when processing pre-fetched records. This change streamlines the data ingestion process from Kafka, leading to more efficient resource utilization.

Highlights

• Performance Optimization in Kafka IO: Removed the computationally expensive Guava Stopwatch from the critical consumer.poll() loop in Kafka IO readers (ReadFromKafkaDoFn and KafkaUnboundedReader).
• Replaced Stopwatch with System.currentTimeMillis(): Switched from Stopwatch (which relies on System.nanoTime()) to a more lightweight System.currentTimeMillis() for tracking remainingTimeout, significantly reducing CPU overhead during rapid Kafka record prefetching.
• Removed Redundant RPC Latency Metrics: Eliminated custom Beam RPC latency metric tracking and associated histograms (KafkaMetrics, KafkaSinkMetrics) as Kafka natively provides this data via JMX (fetch-latency-avg), making the custom implementation redundant and a source of performance overhead.
• Cleaned Up Metric Code and Tests: Removed dead code related to RPC latency metrics from KafkaMetrics.java and KafkaSinkMetrics.java, and updated corresponding test files (KafkaMetricsTest.java, KafkaSinkMetricsTest.java) to reflect these changes.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaMetrics.java
  • Removed updateSuccessfulRpcMetrics method signature from the interface.
  • Removed updateSuccessfulRpcMetrics implementation from NoOpKafkaMetrics.
  • Removed LATENCY_HISTOGRAMS map and perTopicRpcLatencies abstract method.
  • Adjusted create method in KafkaMetricsImpl to no longer initialize perTopicRpcLatencies.
  • Removed the updateSuccessfulRpcMetrics method implementation from KafkaMetricsImpl.
  • Removed the recordRpcLatencyMetrics method.
  • Removed the call to recordRpcLatencyMetrics from flushBufferedMetrics.
• sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaSinkMetrics.java
  • Removed imports for DelegatingHistogram, Histogram, and HistogramData.
  • Removed RPC_LATENCY constant.
  • Removed RpcMethod enum.
  • Removed RPC_METHOD metric label.
  • Removed the createRPCLatencyHistogram method.
• sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaUnboundedReader.java
  • Removed import for Stopwatch.
  • Removed the stopwatch instance variable.
  • Removed stopwatch.start(), stopwatch.stop(), and kafkaResults.updateSuccessfulRpcMetrics calls from the consumerPollLoop method.
• sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaDoFn.java
  • Removed import for Stopwatch.
  • Removed pollTimer (Stopwatch) initialization.
  • Replaced pollTimer.reset().start() with System.currentTimeMillis() for startMillis.
  • Replaced pollTimer.elapsed() with Duration.ofMillis(System.currentTimeMillis() - startMillis) for elapsed calculation.
  • Removed kafkaMetrics.updateSuccessfulRpcMetrics call.
• sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaMetricsTest.java
  • Removed import for containsInAnyOrder.
  • Removed import for Duration.
  • Removed results.updateSuccessfulRpcMetrics call from testNoOpKafkaMetrics.
  • Renamed testKafkaRPCLatencyMetrics to testKafkaBacklogMetrics.
  • Removed all assertions related to RPC latency histograms from testKafkaBacklogMetrics.
  • Removed the testKafkaRPCLatencyMetricsAreNotRecorded test method.
• sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaSinkMetricsTest.java
  • Removed import for Histogram.
  • Removed the testCreatingHistogram method.

Activity

• The code compiled successfully.
• The full Kafka IO test suite was run and passed.
• The timeout logic in the poll loop was verified to function correctly using java.time.Duration and System.currentTimeMillis().

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 by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

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 pull request 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 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. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

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.

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. ↩

[github-actions]

Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment assign set of reviewers

[Abacn]

It's deleting functioning metrics and codes. Please check generated code changes before pushing PRs

[junaiddshaukat]

It's deleting functioning metrics and codes. Please check generated code changes before pushing PRs

Thanks for the feedback, you were right that I accidentally removed existing metrics and related code in my initial change.

I’ve updated the branch to:

• Restore KafkaMetrics, KafkaSinkMetrics, and their tests so no existing metrics or code are deleted anymore.
• Keep the intended optimization for [Task]: Kafka IO: Remove custom latency timer in ReadFromKafkaDoFn in favor Kafka's native metrics #37704, which removes the Guava Stopwatch/System.nanoTime() calls from the hot consumer.poll() loop and uses System.currentTimeMillis() for timeout tracking instead.

I’ve re-run ./gradlew :sdks:java:io:kafka:test locally and it passes.

Please take another look and let me know if you’d like me to narrow or adjust the change further. Thanks

[Abacn]

The TODO is still valid. It's referring to Kafka metrics: https://kafka.apache.org/20/generated/consumer_metrics.html

Stopwatch and System.currentTimeMillis doesn't make a difference here

[junaiddshaukat]

Thanks for the clarification, and I think there was some misunderstanding on my side about the intent of the TODO.

I originally read

“Remove this timer and use the existing fetch‑latency‑avg metric.”

as “the custom Stopwatch‑based timer in the hot consumer.poll() loop is too expensive and should be removed / replaced,” so I focused this PR on changing how we measure the latency (from Stopwatch to System.currentTimeMillis) while keeping the existing Beam metric behavior.

From your comments and the link to the Kafka consumer metrics docs, I now understand that your intent is more about relying on Kafka’s own fetch-latency-avg JMX metric instead of having a separate Beam RpcLatency metric here, rather than just which Java clock is used. In that sense, you’re right: this PR doesn’t actually implement the TODO, it just refactors the timer and doesn’t add much value.

Given that, I’m happy to either:

• Close this PR as a misaligned optimization attempt, or
• If you think it’s worthwhile, follow up with a separate change that truly removes the custom latency timer / RpcLatency metric in favor of Kafka’s metrics (with tests and a clearer design discussion up front).

Please let me know which direction you’d prefer; if you’d rather keep the current behavior and leave the TODO for later, I’ll go ahead and close this PR.