fix(kafka): pause and call poll() immediately on writeAll failure by graytaylor0 · Pull Request #7196 · opensearch-project/data-prepper

graytaylor0 · 2026-10-18T17:13:59Z

Description

When the Kafka consumer buffer writeAll call fails (buffer full or circuit
breaker), the consumer now pauses immediately on the first failure and
calls doPoll() on every retry iteration. This keeps the member in the
group by driving the network I/O that reads JoinGroup/SyncGroup responses
and allows the heartbeat thread to function during rebalances.

Previously, the consumer would retry writeAll for maxRetriesOnException
iterations (calculated as max_poll_interval / 2000) with a 1s sleep
between each retry before pausing. During a rebalance, the heartbeat
thread disables itself after JoinGroup and only re-enables when poll()
processes the JoinGroup response. Without poll() being called, the
coordinator's JoinGroup response sits unread in the socket, and after
session.timeout.ms (60s) expires with no heartbeat or SyncGroup, the
coordinator removes the member. The main thread then gets
UNKNOWN_MEMBER_ID on its next poll, triggering another rebalance.

With this fix, poll() is called on every retry, so JoinGroup/SyncGroup
responses are processed promptly and the heartbeat thread stays active.
Rebalances complete in seconds even when the buffer is full.

Also adds doPoll() in the pauseConsumePredicate path so that external
pause conditions (like circuit breaker) also maintain group membership.

Issues Resolved

Resolves #[Issue number to be closed when this PR is merged]

Check List

New functionality includes testing.
New functionality has a documentation issue. Please link to it in this PR.
- New functionality has javadoc added
Commits are signed with a real name per the DCO

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

github-actions · 2026-10-18T17:15:09Z

PR Reviewer Guide 🔍

(Review updated until commit `15fb76f`)

Here are some key observations to aid the review process:

🧪 PR contains tests
🔒 No security concerns identified
✅ No TODO sections
🔀 No multiple PR themes
⚡ Recommended focus areas for review Unexpected records lost on paused poll When `doPoll()` returns records while the consumer is paused (which can happen immediately after `consumer.pause()` was called but before the broker-side state fully reflects the pause, or on first failure after records were just prefetched), the code logs a warning, seeks back, breaks out of the retry loop, and returns from `processRecords` without ever retrying the `buffer.writeAll` of the current `eventRecords`. Those records are effectively dropped (not written to buffer, not acknowledged, and the offset is reset so they'll be re-polled — which is fine for at-least-once, but combined with removing the retry counter, the first iteration after failure now always calls `doPoll()`, significantly increasing the chance of hitting this path and churning on the same batch instead of actually draining to the buffer when space becomes available. if (paused && !shutdownInProgress.get()) { ConsumerRecords<String, ?> records = doPoll(); if (records.count() > 0) { LOG.warn("Unexpected records received while the consumer is paused. Resetting the partitions to retry from last read pointer"); synchronized(this) { partitionsToReset.addAll(consumer.assignment()); }; break; } } Pause predicate path never resumes In the new `pauseConsumePredicate.pauseConsuming()` branch, `paused = true` and `consumer.pause(consumer.assignment())` are set, then `continue` jumps back to the top of the loop. If the predicate later returns false, control falls through to the `else if(paused)` branch which resumes — but only if no exception caused re-entry. More importantly, `doPoll()` is called while paused and its returned records are discarded (unlike in `processRecords`, there's no seek-back), so any records returned by `poll()` here while assignments are not yet fully paused would be silently lost. } else if (pauseConsumePredicate.pauseConsuming()) { LOG.debug("Pause and skip consuming from Kafka topic due to an external condition: {}", pauseConsumePredicate); paused = true; consumer.pause(consumer.assignment()); doPoll(); Thread.sleep(1000); continue;

github-actions · 2026-10-18T17:15:30Z

PR Code Suggestions ✨

Latest suggestions up to 15fb76f

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Avoid dropping records returned by poll Calling `doPoll()` on paused partitions can still return records if partitions were just assigned/revoked, but more importantly its result is discarded here. If records are returned (e.g., from a pre-pause fetch buffer), they will be silently dropped causing data loss. Capture the result and, if non-empty, seek back to the last committed offsets like the `processRecords` retry path does. data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumer.java [429] } else if (pauseConsumePredicate.pauseConsuming()) { LOG.debug("Pause and skip consuming from Kafka topic due to an external condition: {}", pauseConsumePredicate); paused = true; consumer.pause(consumer.assignment()); - doPoll(); + ConsumerRecords<String, ?> records = doPoll(); + if (records.count() > 0) { + LOG.warn("Unexpected records received while the consumer is paused. Resetting the partitions to retry from last read pointer"); + synchronized(this) { + partitionsToReset.addAll(records.partitions()); + } + } Thread.sleep(1000); continue; Suggestion importance[1-10]: 7 __ Why: Valid concern: discarding the return value of `doPoll()` while paused could lose records that arrive before pause takes effect. The suggestion mirrors the handling pattern used in `processRecords`, improving correctness.	Medium
General	Reconsider immediate pause on first failure Removing the retry threshold means the consumer now pauses immediately on the very first transient exception (including a brief `TimeoutException` from `writeAll`). This may cause unnecessary pause/resume churn and extra `poll()` round trips for short-lived buffer pressure. Consider keeping a small retry count before pausing, or document that this is the intended behavior change. data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumer.java [536-539] +} catch (Exception e) { + if (!paused) { + paused = true; + consumer.pause(consumer.assignment()); + } - Suggestion importance[1-10]: 2 __ Why: The `existing_code` and `improved_code` are identical, so no concrete change is proposed; it's merely a design discussion comment.	Low

Previous suggestions

Suggestions up to commit 33fe153

Category	Suggestion	Impact
Possible issue	Avoid discarding records returned after pause Calling `doPoll()` right after pausing will still return any records that were already buffered client-side before the pause, which will be silently discarded here (unlike in `processRecords`, there's no seek-back on the partitions). This can cause data loss when the external pause predicate triggers. Consider seeking back on any partitions whose records were returned, matching the handling in `processRecords`. data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumer.java [425-431] } else if (pauseConsumePredicate.pauseConsuming()) { LOG.debug("Pause and skip consuming from Kafka topic due to an external condition: {}", pauseConsumePredicate); paused = true; consumer.pause(consumer.assignment()); - doPoll(); + ConsumerRecords<String, ?> records = doPoll(); + if (records.count() > 0) { + LOG.warn("Unexpected records received while the consumer is paused. Resetting the partitions to retry from last read pointer"); + synchronized(this) { + records.partitions().forEach(partition -> { + long lastOffset = records.records(partition).get(0).offset(); + consumer.seek(partition, lastOffset); + }); + } + } Thread.sleep(1000); continue; Suggestion importance[1-10]: 8 __ Why: Valid concern - calling `doPoll()` without handling returned records could cause data loss as records fetched before pause would be silently discarded, unlike the handling in `processRecords`.	Medium
General	Allow retry loop to exit on shutdown Removing `maxRetriesOnException` means `processRecords` can now loop indefinitely on repeated `writeAll` failures with no upper bound. Even though the consumer is paused and polled to keep the session alive, consider adding a safety limit or shutdown-aware exit so a persistently failing buffer does not block the thread forever; also check `shutdownInProgress` in the while condition. data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumer.java [535-539] } catch (Exception e) { if (!paused) { paused = true; consumer.pause(consumer.assignment()); } + if (shutdownInProgress.get()) { + break; + } Suggestion importance[1-10]: 7 __ Why: Reasonable concern since removing `maxRetriesOnException` creates an unbounded retry loop; adding a `shutdownInProgress` check prevents the thread from blocking indefinitely during shutdown.	Medium

Suggestions up to commit d920cd9

Category	Suggestion	Impact
Possible issue	Only poll during retries when paused `doPoll()` is now invoked on every retry regardless of whether the consumer was paused. Before pausing, the partitions are not paused, so `doPoll()` will return new records and trigger the "Unexpected records received while the consumer is paused" branch and partition seek-reset even in the normal (not-yet-paused) retry path. Guard the poll with `if (paused && !shutdownInProgress.get())` to preserve the original intent. data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumer.java [546-551] Thread.sleep(RETRY_ON_EXCEPTION_SLEEP_MS); -if (!shutdownInProgress.get()) { +if (paused && !shutdownInProgress.get()) { ConsumerRecords<String, ?> records = doPoll(); if (records.count() > 0) { LOG.warn("Unexpected records received while the consumer is paused. Resetting the partitions to retry from last read pointer"); synchronized(this) { Suggestion importance[1-10]: 8 __ Why: In the new code, the consumer is now paused immediately on the first exception (via `if (!paused)`), so `paused` will typically be true during retries. However, the suggestion correctly identifies that the original check `if (paused)` was intentional and the change to `if (!shutdownInProgress.get())` removes this guard, which could cause unexpected behavior if the exception path somehow runs without pausing. This is a valid correctness concern.	Medium

Suggestions up to commit 636272d

Category	Suggestion	Impact
Possible issue	Only poll during retries when paused `doPoll()` is now invoked unconditionally on every retry regardless of whether the consumer was actually paused. If `paused` is false (e.g., writeAll failed but pause path was skipped), polling here may fetch new records that are then unexpectedly reset via `seek`, causing data re-processing issues. Guard the poll with `if (paused &&` `!shutdownInProgress.get())` to preserve the original intent. data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumer.java [547-552] Thread.sleep(RETRY_ON_EXCEPTION_SLEEP_MS); -if (!shutdownInProgress.get()) { +if (paused && !shutdownInProgress.get()) { ConsumerRecords<String, ?> records = doPoll(); if (records.count() > 0) { LOG.warn("Unexpected records received while the consumer is paused. Resetting the partitions to retry from last read pointer"); synchronized(this) { Suggestion importance[1-10]: 7 __ Why: Valid concern: the old code only polled when `paused` was true, and the new code polls regardless, which could fetch records that are then reset via `seek`. However, in the new code `paused` is always set to true before this path is reached, so the practical impact may be limited.	Medium
Possible issue	Resume consumer when external pause ends After pausing and polling here, `continue` is executed without ever resuming the consumer when `pauseConsumePredicate.pauseConsuming()` later returns false. The next iteration goes through `consumeRecords()` which polls while partitions are still paused, causing permanent starvation. Resume the assignment (and reset `paused`) when the predicate no longer requires pausing. data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/consumer/KafkaCustomConsumer.java [425-431] } else if (pauseConsumePredicate.pauseConsuming()) { LOG.debug("Pause and skip consuming from Kafka topic due to an external condition: {}", pauseConsumePredicate); paused = true; consumer.pause(consumer.assignment()); doPoll(); Thread.sleep(1000); continue; +} else if (paused) { + LOG.debug("Resume consuming from Kafka topic."); + consumer.resume(consumer.assignment()); + paused = false; Suggestion importance[1-10]: 7 __ Why: Valid observation that there appears to be no resume path shown here when the predicate stops requiring pause. However, the diff shows an `else if(paused)` branch at line 432 that likely handles resume, so this may already be covered in surrounding code.	Medium

github-actions · 2026-10-18T17:29:27Z

Persistent review updated to latest commit d920cd9

github-actions · 2026-10-18T19:16:24Z

Persistent review updated to latest commit 33fe153

dlvenable · 2026-10-19T13:57:43Z


+    @ParameterizedTest
+    @MethodSource("provideExceptionsFromBufferWrite")
+    public void testWriteAllFailurePausesImmediatelyAndCallsPoll(final Exception bufferException) throws Exception {


Do we have existing unit tests that we can add an assertion to on the poll?

When the Kafka consumer buffer writeAll call fails (buffer full or circuit breaker), the consumer now pauses immediately on the first failure and calls doPoll() on every retry iteration. This keeps the member in the group by driving the network I/O that reads JoinGroup/SyncGroup responses and allows the heartbeat thread to function during rebalances. Previously, the consumer would retry writeAll for maxRetriesOnException iterations (calculated as max_poll_interval / 2000) with a 1s sleep between each retry before pausing. During a rebalance, the heartbeat thread disables itself after JoinGroup and only re-enables when poll() processes the JoinGroup response. Without poll() being called, the coordinator's JoinGroup response sits unread in the socket, and after session.timeout.ms (60s) expires with no heartbeat or SyncGroup, the coordinator removes the member. The main thread then gets UNKNOWN_MEMBER_ID on its next poll, triggering another rebalance. With this fix, poll() is called on every retry, so JoinGroup/SyncGroup responses are processed promptly and the heartbeat thread stays active. Rebalances complete in seconds even when the buffer is full. Also adds doPoll() in the pauseConsumePredicate path so that external pause conditions (like circuit breaker) also maintain group membership. Signed-off-by: Taylor Gray <tylgry@amazon.com>

github-actions · 2026-10-19T16:24:50Z

Persistent review updated to latest commit 15fb76f

graytaylor0 requested review from KarstenSchnitter, Zhangxunmt, dinujoh, divbok, dlvenable, kkondaka, oeyh, san81, sb2k16, srikanthjg and srikanthpadakanti as code owners October 18, 2026 17:14

graytaylor0 force-pushed the KafkaBufferPollFix branch from 636272d to d920cd9 Compare October 18, 2026 17:28

graytaylor0 force-pushed the KafkaBufferPollFix branch from d920cd9 to 33fe153 Compare October 18, 2026 19:15

dlvenable reviewed Oct 19, 2026

View reviewed changes

graytaylor0 force-pushed the KafkaBufferPollFix branch from 33fe153 to 15fb76f Compare October 19, 2026 16:23

dlvenable approved these changes Oct 19, 2026

View reviewed changes

graytaylor0 merged commit a81b295 into opensearch-project:main Oct 19, 2026
74 of 77 checks passed

Navigation Menu

Repository navigation

fix(kafka): pause and call poll() immediately on writeAll failure - #7196

fix(kafka): pause and call poll() immediately on writeAll failure#7196
graytaylor0 merged 1 commit into
opensearch-project:mainfrom
graytaylor0:KafkaBufferPollFix

graytaylor0 commented Oct 18, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Oct 18, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Oct 18, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Oct 18, 2026

Uh oh!

github-actions Bot commented Oct 18, 2026

Uh oh!

dlvenable Oct 19, 2026

Uh oh!

github-actions Bot commented Oct 19, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

graytaylor0 commented Oct 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Issues Resolved

Check List

Uh oh!

github-actions Bot commented Oct 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR Reviewer Guide 🔍

(Review updated until commit 15fb76f)

Uh oh!

github-actions Bot commented Oct 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR Code Suggestions ✨

Previous suggestions

Uh oh!

github-actions Bot commented Oct 18, 2026

Uh oh!

github-actions Bot commented Oct 18, 2026

Uh oh!

dlvenable Oct 19, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Oct 19, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

graytaylor0 commented Oct 18, 2026 •

edited

Loading

github-actions Bot commented Oct 18, 2026 •

edited

Loading

(Review updated until commit `15fb76f`)

github-actions Bot commented Oct 18, 2026 •

edited

Loading