[KafkaIO] Use consumer position and lag to estimate end offsets

[sjvanrossum]

Use Consumer.position(TopicPartition) and Consumer.currentLag(TopicPartition) to update the end offset estimator (reduced to an AtomicLong per restriction tracker) using local state.

This change reuses local metadata from recently completed record fetches instead of periodically requesting the broker for end offsets.

This PR was split from #36834 and depends on the changes in #39284 and #39297.

---

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, 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 optimizes Kafka end offset estimation by leveraging local consumer state, which reduces network overhead by avoiding frequent broker requests. Additionally, the PR upgrades the Kafka client to version 3.9.2 and modernizes test infrastructure by replacing deprecated Kafka server utilities. These changes also include a cleanup of legacy build configurations to maintain a cleaner project structure.

Highlights

• Offset Estimation Optimization: Refactored the end offset estimation in ReadFromKafkaDoFn to use local consumer position and lag, eliminating the need for periodic broker requests.
• Kafka Version Upgrade: Upgraded the project's Kafka dependency to version 3.9.2.
• Test Utility Modernization: Replaced the deprecated KafkaServerStartable with KafkaServer in test utilities and updated associated configurations.
• Build Cleanup: Removed legacy Kafka version configurations and cleaned up project inclusions in build files.

New Features

🧠 You can now enable Memory (public preview) 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.

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 the 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 counterproductive. 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.

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

[gemini-code-assist]

Code Review

This pull request upgrades the Kafka dependency to version 3.9.2, cleans up older Kafka integration test modules, and refactors the offset estimation in ReadFromKafkaDoFn to lazily track partition lag using the polling consumer instead of a dedicated estimator consumer. Feedback on these changes highlights critical issues with using Long.MIN_VALUE as the default uninitialized offset, which can cause incorrect progress tracking (reporting zero remaining work) and negative backlog calculations that disrupt autoscaling. Additionally, the reviewer recommends logging exceptions in the empty catch block during polling and avoiding fragile Scala synthetic methods when instantiating KafkaServer in tests.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Since latestOffsetEstimator is initialized to Long.MIN_VALUE, the GrowableOffsetRangeTracker will receive Long.MIN_VALUE as the end offset estimate before the first poll completes. This causes getProgress() to cap the end offset to the current position and report 0 remaining work, which can incorrectly signal to the runner that the partition is fully processed. Returning Long.MAX_VALUE when the estimate is unknown (i.e., Long.MIN_VALUE) represents an infinite/unknown end offset, which is much safer and more appropriate for an unbounded source.

    final AtomicLong latestOffsetEstimator =
        latestOffsetEstimatorCache.getUnchecked(kafkaSourceDescriptor);
    return new UnsplittableRestrictionTracker<>(
        new GrowableOffsetRangeTracker(
            restriction.getFrom(),
            () -> {
              long estimate = latestOffsetEstimator.get();
              return estimate == Long.MIN_VALUE ? Long.MAX_VALUE : estimate;
            }));

[sjvanrossum]

This explanation and suggestion look wrong to me. Provide a source reference for the runner behavior.
Docs for GrowableOffsetRangeTracker.RangeEndEstimator state:

Return {@code Long.MIN_VALUE} if an estimate can not be provided.

[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

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 58.04%. Comparing base (7d682e2) to head (9ffbc88).
⚠️ Report is 6 commits behind head on master.

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #39285      +/-   ##
============================================
+ Coverage     54.76%   58.04%   +3.28%     
- Complexity     1716    13052   +11336     
============================================
  Files          1066     2515    +1449     
  Lines        169067   263741   +94674     
  Branches       1255    10761    +9506     
============================================
+ Hits          92587   153090   +60503     
- Misses        74263   104904   +30641     
- Partials       2217     5747    +3530     

Flag	Coverage Δ
java	64.23% <ø> (-3.38%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.