[KafkaIO] Fix off-by-one message backlog in SDF read

[udayaw]

Problem

ReadFromKafkaDoFn reports a fully caught-up partition as having a backlog of
1, never 0. The cause is a collision between two offset conventions.

Kafka's offsets are exclusive; the tracker's claimed offset is not

Consumer#position() is documented as "the offset of the next record that
will be fetched" — it is exclusive by design, not off by one. currentLag() is
logEndOffset - position(). So the estimate KafkaIO installs in the tracker,

final long position = consumer.position(topicPartition);
consumer.currentLag(topicPartition)
    .ifPresent(lag -> latestOffsetEstimator.lazySet(position + lag));

is exactly the log end offset — exclusive. That is precisely what the tracker
asks for; GrowableOffsetRangeTracker.RangeEndEstimator states "The end offset
is exclusive for the range." The estimator is correct.

lastAttemptedOffset, however, is the last inclusive offset claimed. And
getProgress() subtracts one from the other without converting:

final long completedEnd = lastAttemptedOffset == null ? range.getFrom() : lastAttemptedOffset;
final long remainingEnd = Math.max(completedEnd, rangeEndEstimator.estimate());
return Progress.from(completedEnd - range.getFrom(), remainingEnd - completedEnd);

Once at least one record has been claimed, workRemaining == realBacklog + 1
(and workCompleted is one short). Worked through, for a partition holding
offsets 100–104 that the reader has fully drained:

value
log end offset	105	exclusive
position()	105	exclusive — next record to fetch
currentLag()	0	reader is caught up
estimate (position + lag)	105	exclusive ✅
lastAttemptedOffset	104	inclusive
workRemaining	105 - 104 = 1	❌ should be 0

ReadFromKafkaDoFn then feeds that straight to the timestamp policy:

new TimestampPolicyContext(
    (long) ((HasProgress) tracker).getProgress().getWorkRemaining(), Instant.now());

Why it matters

Both CustomTimestampPolicyWithLimitedDelay and
TimestampPolicyFactory.LogAppendTimePolicy gate their idle-advance branch on
ctx.getMessageBacklog() == 0. Because that value can never reach zero, the
branch is unreachable, and an idle partition's watermark stays pinned at
lastRecordTimestamp - maxDelay for as long as the partition stays quiet.
Downstream event-time windows and triggers never fire.

Only the SDF path is affected. KafkaUnboundedReader.backlogMessageCount()
computes latestOffset - nextOffset — exclusive minus exclusive — and does reach
0. Dataflow always selects the SDF path (KafkaIO.Read#runnerPrefersLegacyRead
returns false for any org.apache.beam.runners.dataflow.* runner) unless the
use_deprecated_read experiment is set, so pipelines there are exposed by
default.

Observed in production on a Dataflow streaming job: a topic stopped producing and
the job reported data freshness growing 1:1 with wall clock for over eight hours,
while committed offsets showed the reader had been fully caught up the whole
time. It simply had no way to say so.

Evidence this is an oversight rather than intent

Two places in the existing code handle the exclusive/inclusive distinction
correctly, one of them three lines from the defect:

1. The claim path converts explicitly. ReadFromKafkaDoFn subtracts one from
   position() to obtain a claimable inclusive offset:

   if (expectedOffset < (expectedOffset = consumer.position(topicPartition))) {
     if (!tracker.tryClaim(expectedOffset - 1)) {

   The code knows the two conventions differ. getProgress() just never converts
   back.

2. The backlogBytes gauge gets it right, ~20 lines below
   updateWatermarkManually in the same method, using the exclusive
   expectedOffset:

   BigDecimal.valueOf(Math.max(expectedOffset, latestOffsetEstimator.get()))
       .subtract(BigDecimal.valueOf(expectedOffset), MathContext.DECIMAL128)

   Same method, two conventions, and the metric uses the correct one.

Notably, every offset in processElement is already exclusive-next —
tracker.currentRestriction().getFrom(), rawRecord.offset() + 1, and
consumer.position(). lastAttemptedOffset, reachable only through the tracker,
is the sole inclusive value in the picture.

Change

updateWatermarkManually now takes the exclusive next offset and the latest end
offset estimate directly, and computes
max(estimatedEndOffset, nextOffset) - nextOffset — exclusive minus exclusive,
matching both the gauge and the legacy reader. Every call site already had
expectedOffset in hand, so no new state is threaded through.

TimestampPolicy.PartitionContext#getMessageBacklog's javadoc previously read
"latest offset of the partition - last processed record offset", which
describes the buggy inclusive arithmetic and is plausibly how the slip survived
review. Reworded to state that both offsets are exclusive and that zero means
fully caught up.

Long.MIN_VALUE — the sentinel latestOffsetEstimator holds when the position is
undefined or out of range — still yields a backlog of 0 via the max, matching
today's behaviour on that path.

Compacted topics

The value this method reports is a count of offsets, not of messages —
compaction deletes records while leaving the offset span intact, so on a
compacted topic the backlog overstates how many records actually remain. That is
pre-existing and unchanged here; getSize already notes it ("Compacted topics
may hold less records than the estimated offset range due to record deletion
within a partition"), and the legacy KafkaUnboundedReader has the same
characteristic since it also computes latestOffset - nextOffset. This change
removes a systematic off-by-one that affects every topic, compacted or not.

The idle-advance path does still reach zero on a compacted topic. The
"non-visible progress" branch claims up to consumer.position() - 1 whenever the
position advances without any records being returned, which is exactly what a
fetch across a compacted region looks like, so expectedOffset follows the
consumer past the gaps. For a partition whose surviving records are at offsets
0, 2 and 4 with a log end offset of 5, the backlog after the final record is 0
with this change and 1 without it.

What this change does not address is a consumer position that stalls below the
log end offset with nothing left to fetch; the backlog would stay positive and
the watermark would not advance. That is a fetch-position concern rather than an
arithmetic one, is unaffected by this change, and behaves identically on the
legacy read path.

Alternative considered: converting inside the tracker

The root cause is in GrowableOffsetRangeTracker.getProgress(), and
completedEnd = lastAttemptedOffset + 1 would fix workCompleted too. I did not
do that here because:

• it changes progress values for every SDF using the tracker, and runners
  consume those for splitting and autoscaling decisions;
• it needs an overflow guard for the lastAttemptedOffset == Long.MAX_VALUE
  (range done) case, since getProgress feeds the subtraction through
  UnsignedLong.fromLongBits.

The two fixes are independent — this one derives the backlog without consulting
the tracker — so a later tracker fix will not double-correct. Happy to switch to
it, or file it as a follow-up, if reviewers prefer.

Scope note: getSize() is not affected. It builds a fresh tracker per
call, so lastAttemptedOffset is null, completedEnd == range.getFrom(), and
workRemaining == estimate - from — correct, because restriction starts are
exclusive-next offsets. The in-flight getProgress() the runner queries during
bundle processing does stay off by one until the tracker itself is fixed; that
affects splitting and progress reporting, not watermarks.

Tests

Four tests in ReadFromKafkaDoFnTest, all driving the real
GrowableOffsetRangeTracker through restrictionTracker():

• testMessageBacklogReachesZeroWhenPartitionIsCaughtUp — exact backlog sequence
  [2, 1, 0, 0] for three records on a partition with end offset 3.
• testMessageBacklogExcludesRecordsAlreadyRead — [7, 6, 5, 5] with five
  records left unread.
• testWatermarkAdvancesForIdleCaughtUpPartition — a real
  CustomTimestampPolicyWithLimitedDelay advances past the last record's
  timestamp once the partition drains.
• testWatermarkDoesNotAdvanceForIdlePartitionWithBacklog — the same policy
  stays at lastRecordTimestamp - maxDelay while records remain.

All four fail before the change and pass after. They use a new
IdlingMockKafkaConsumer, which returns one batch then empty polls and tracks
position so currentLag is meaningful; the existing SimpleMockKafkaConsumer
returns a fixed position and cannot express "caught up".

---

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.

[github-actions]

Assigning reviewers:

R: @chamikaramj for label java.
R: @johnjcasey for label kafka.

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

• stop reviewer notifications - opt out of the automated review tooling
• remind me after tests pass - tag the comment author after tests pass
• waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).