fix(kafka source): wait for aborted partition consumers before continuing a rebalance

[sandervandegeijn]

Summary

After a consumer group rebalance the kafka source could silently stop consuming some of its partitions. Lag and memory grew, nothing was logged, and only a restart helped. This is the behaviour reported in #22006.

Root cause

Every partition is consumed through its own rdkafka StreamPartitionQueue. All queues split for the same partition wrap one librdkafka fetch queue. split_partition_queue installs a wake-up callback on that queue and Drop for StreamPartitionQueue removes it unconditionally, so the last writer wins.

When pending acknowledgements could not be drained within drain_timeout_ms during a revoke (for example under sink backpressure), the coordinator aborted the partition tasks and immediately released the rebalance callback. Tokio drops an aborted task asynchronously. If the partition was re-assigned to the same consumer before the old task was actually dropped, the new queue was split first and the old drop then removed the new queue's wake-up. The new task was never woken again and fetched messages piled up unread.

Fix

• The coordinator stays in the draining state after aborting until JoinSet::join_next_with_id has reaped every aborted task. Tokio drops the task's future before yielding its result, so the old queue is gone before the rebalance callback is released.
• Task ids guard the abort handle and end signal bookkeeping, so a task that finishes late can no longer remove the entries of a newer task for the same partition. Previously this could also leave a partition undrained on the next revoke.
• Aborted or panicked tasks are now handled explicitly instead of being silently dropped by the select! pattern, so the drain can no longer hang on them.
• The wait for aborted tasks is not bounded, because giving up would re-open the race. It is logged at warn every drain_timeout_ms while it lasts. In practice it ends at the task's next .await.
• Aborted tasks no longer trigger a synchronous commit each. The client commits once more after the drain anyway, which covers their stored offsets.

This covers the path where partitions revoked under backpressure come back to the same consumer. Other reports in the #22006 thread (partition count changes, very large consumer groups) are not addressed here.

I developed and verified this fix together with Claude Code. I reviewed every line of it and the reasoning above is my own understanding of the problem.

References

Closes #22006

Vector configuration

The integration tests use a source equivalent to this, with a deliberately short drain timeout so the deadline fires during the rebalance:

sources:
  kafka:
    type: kafka
    bootstrap_servers: localhost:9091
    group_id: test-group
    topics: [test-topic]
    auto_offset_reset: beginning
    session_timeout_ms: 6000
    drain_timeout_ms: 100
    librdkafka_options:
      partition.assignment.strategy: range,roundrobin   # and cooperative-sticky

How did you test this PR?

Reproduction and proof. I confirmed the mechanism in the rdkafka 0.39.0 source and then reproduced the stall deterministically against a live broker. With a scratch-only patch that makes every message take a long synchronous decode, the partition tasks are mid-decode when the drain deadline fires and their abort lands late. A second blocked group member joins and leaves twice. I ran this scenario from the same branch with and without the new waiting logic:

Variant	range,roundrobin	cooperative-sticky
Waiting logic reverted (old behaviour)	3/3 runs stalled, two partitions never received follow-up messages	3/3 pass
This PR	3/3 pass	3/3 pass

Cooperative-sticky does not stall in this scenario because revoked partitions move to the other member instead of returning to the same consumer immediately. The reproduction test itself is not part of this PR because it needs a test-only hook in production code.

Integration tests. Two new tests (consumes_after_aborted_drain_*) exercise the aborted-drain re-assignment path under backpressure for both rebalance protocols. They passed in 9 out of 9 runs. With tracing enabled they show the drain deadline firing and the tasks being aborted, and never the "still waiting" warning. The full kafka-integration-tests suite passes on this branch (the two sink SASL/TLS tests only fail locally because my scratch broker has no SASL or TLS listener).

Production. I have been running this fix in production for a week. The stall has not recurred and I have seen no side effects.

Static checks. make check-clippy (kafka features), make check-fmt, cargo vdev check changelog-fragments, and markdownlint-cli2 on the changelog fragment are clean. The generated docs in website/cue are updated to match the changed drain_timeout_ms description.

Does this PR include user facing changes?

• Yes. Please add a changelog fragment based on our guidelines.
• No. A maintainer will apply the no-changelog label to this PR.

The user-visible change is that a rebalance or shutdown can take slightly longer than drain_timeout_ms when partitions did not drain in time, and the drain_timeout_ms documentation now states that undrained partitions are stopped.

Contributor Guidelines

• Please read our Vector contributor resources.
• Do not hesitate to use @vectordotdev/vector to reach out to us regarding this PR.
• Before pushing, follow our pre-push guidance.
• After a review is requested, please avoid force pushes to help us review incrementally.
  • Feel free to push as many commits as you want. They will be squashed into one before merging.
  • For example, you can run git merge origin master and git push.

[github-actions]

All contributors have signed the CLA ✍️ ✅
_{Posted by the CLA Assistant Lite bot.}

[sandervandegeijn]

I have read the CLA Document and I hereby sign the CLA