KAFKA-14220: Update WorkerSinkTask.java

[ghost]

No description provided.

[C0urante]

Thanks @kumarpritam863, this is an important fix. I did some research on the history behind incremental consumer rebalancing and how it relates to the rebalance listener API, and found that this PR contained some great details, including an explicit confirmation on the expected ordering of callback invocation:

The ordering of the callback would be the following:
a. Callback onPartitionsRevoked / onPartitionsLost triggered.
b. Update the assignment (both revoked and added).
c. Callback onPartitionsAssigned triggered.

One case I'm wondering about is when the task fails, which would cause its consumer to leave the group. I believe that this would cause onPartitionsRevoked to be invoked in the rebalance listener, but not onPartitionsAssigned, so the metrics for the task (which are still available for failed tasks until they are explicitly revoked from the worker) would become inaccurate.

Could we forcibly reset the partition count metric to zero in WorkerSinkTask::close to address this case? (Probably with a brief comment explaining why we have this special logic).

It'd also be nice to see some tests for this logic, both to verify correctness and help prevent regression.