/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.kafka.clients.consumer;


import org.apache.kafka.clients.ClientsTestUtils;
import org.apache.kafka.clients.ClientsTestUtils.TestConsumerReassignmentListener;
import org.apache.kafka.clients.producer.Producer;
import org.apache.kafka.common.MetricName;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.test.ClusterInstance;
import org.apache.kafka.common.test.api.ClusterConfigProperty;
import org.apache.kafka.common.test.api.ClusterTest;
import org.apache.kafka.common.test.api.ClusterTestDefaults;
import org.apache.kafka.common.test.api.Type;
import org.apache.kafka.common.utils.Utils;
import org.apache.kafka.test.TestUtils;

import org.junit.jupiter.api.BeforeEach;

import java.time.Duration;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashSet;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;
import java.util.stream.IntStream;

import static org.apache.kafka.clients.ClientsTestUtils.awaitRebalance;
import static org.apache.kafka.clients.ClientsTestUtils.consumeAndVerifyRecords;
import static org.apache.kafka.clients.ClientsTestUtils.ensureNoRebalance;
import static org.apache.kafka.clients.ClientsTestUtils.sendRecords;
import static org.apache.kafka.clients.ClientsTestUtils.waitForPollThrowException;
import static org.apache.kafka.clients.CommonClientConfigs.MAX_POLL_INTERVAL_MS_CONFIG;
import static org.apache.kafka.clients.CommonClientConfigs.SESSION_TIMEOUT_MS_CONFIG;
import static org.apache.kafka.clients.consumer.ConsumerConfig.AUTO_OFFSET_RESET_CONFIG;
import static org.apache.kafka.clients.consumer.ConsumerConfig.CLIENT_ID_CONFIG;
import static org.apache.kafka.clients.consumer.ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG;
import static org.apache.kafka.clients.consumer.ConsumerConfig.GROUP_ID_CONFIG;
import static org.apache.kafka.clients.consumer.ConsumerConfig.GROUP_PROTOCOL_CONFIG;
import static org.apache.kafka.clients.consumer.ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG;
import static org.apache.kafka.clients.consumer.ConsumerConfig.MAX_POLL_RECORDS_CONFIG;
import static org.apache.kafka.coordinator.group.GroupCoordinatorConfig.CONSUMER_GROUP_HEARTBEAT_INTERVAL_MS_CONFIG;
import static org.apache.kafka.coordinator.group.GroupCoordinatorConfig.CONSUMER_GROUP_MIN_HEARTBEAT_INTERVAL_MS_CONFIG;
import static org.apache.kafka.coordinator.group.GroupCoordinatorConfig.GROUP_INITIAL_REBALANCE_DELAY_MS_CONFIG;
import static org.apache.kafka.coordinator.group.GroupCoordinatorConfig.GROUP_MIN_SESSION_TIMEOUT_MS_CONFIG;
import static org.apache.kafka.coordinator.group.GroupCoordinatorConfig.OFFSETS_TOPIC_PARTITIONS_CONFIG;
import static org.apache.kafka.coordinator.group.GroupCoordinatorConfig.OFFSETS_TOPIC_REPLICATION_FACTOR_CONFIG;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertTrue;

@ClusterTestDefaults(
    types = {Type.KRAFT},
    brokers = PlaintextConsumerPollTest.BROKER_COUNT,
    serverProperties = {
        @ClusterConfigProperty(key = OFFSETS_TOPIC_PARTITIONS_CONFIG, value = "1"),
        @ClusterConfigProperty(key = OFFSETS_TOPIC_REPLICATION_FACTOR_CONFIG, value = "3"),
        @ClusterConfigProperty(key = GROUP_MIN_SESSION_TIMEOUT_MS_CONFIG, value = "100"),
        @ClusterConfigProperty(key = CONSUMER_GROUP_HEARTBEAT_INTERVAL_MS_CONFIG, value = "500"),
        @ClusterConfigProperty(key = CONSUMER_GROUP_MIN_HEARTBEAT_INTERVAL_MS_CONFIG, value = "500"),
        @ClusterConfigProperty(key = GROUP_INITIAL_REBALANCE_DELAY_MS_CONFIG, value = "10"),
    }
)
public class PlaintextConsumerPollTest {

    public static final int BROKER_COUNT = 3;
    public static final double EPSILON = 0.1;
    public static final long GROUP_MAX_SESSION_TIMEOUT_MS = 60000L;
    private final ClusterInstance cluster;
    private final String topic = "topic";
    private final TopicPartition tp = new TopicPartition(topic, 0);
    private final TopicPartition tp2 = new TopicPartition(topic, 1);

    public PlaintextConsumerPollTest(ClusterInstance cluster) {
        this.cluster = cluster;
    }

    @BeforeEach
    public void setup() throws InterruptedException {
        cluster.createTopic(topic, 2, (short) BROKER_COUNT);
    }

    @ClusterTest
    public void testClassicConsumerMaxPollRecords() throws InterruptedException {
        testMaxPollRecords(GroupProtocol.CLASSIC);
    }

    @ClusterTest
    public void testAsyncConsumerMaxPollRecords() throws InterruptedException {
        testMaxPollRecords(GroupProtocol.CONSUMER);
    }

    private void testMaxPollRecords(GroupProtocol groupProtocol) throws InterruptedException {
        var maxPollRecords = 100;
        var numRecords = 5000;
        Map<String, Object> config = Map.of(
            MAX_POLL_RECORDS_CONFIG, maxPollRecords,
            GROUP_PROTOCOL_CONFIG, groupProtocol.name().toLowerCase(Locale.ROOT)
        );
        var startingTimestamp = System.currentTimeMillis();
        sendRecords(cluster, tp, numRecords, startingTimestamp);
        try (Consumer<byte[], byte[]> consumer = cluster.consumer(config)) {
            consumer.assign(List.of(tp));
            consumeAndVerifyRecords(
                consumer,
                tp,
                numRecords,
                maxPollRecords,
                0,
                0,
                startingTimestamp,
                -1
            );
        }
    }

    @ClusterTest
    public void testClassicConsumerMaxPollIntervalMs() throws InterruptedException {
        testMaxPollIntervalMs(Map.of(
            MAX_POLL_INTERVAL_MS_CONFIG, 1000,
            GROUP_PROTOCOL_CONFIG, GroupProtocol.CLASSIC.name().toLowerCase(Locale.ROOT),
            HEARTBEAT_INTERVAL_MS_CONFIG, 500,
            SESSION_TIMEOUT_MS_CONFIG, 2000
        ));
    }

    @ClusterTest
    public void testAsyncConsumerMaxPollIntervalMs() throws InterruptedException {
        testMaxPollIntervalMs(Map.of(
            MAX_POLL_INTERVAL_MS_CONFIG, 1000,
            GROUP_PROTOCOL_CONFIG, GroupProtocol.CONSUMER.name().toLowerCase(Locale.ROOT)
        ));
    }

    private void testMaxPollIntervalMs(Map<String, Object> config) throws InterruptedException {
        try (Consumer<byte[], byte[]> consumer = cluster.consumer(config)) {
            var listener = new TestConsumerReassignmentListener();
            consumer.subscribe(List.of(topic), listener);

            // rebalance to get the initial assignment
            awaitRebalance(consumer, listener);
            assertEquals(1, listener.callsToAssigned);
            assertEquals(0, listener.callsToRevoked);

            // after we extend longer than max.poll a rebalance should be triggered
            // NOTE we need to have a relatively much larger value than max.poll to let heartbeat expired for sure
            TimeUnit.MILLISECONDS.sleep(3000);

            awaitRebalance(consumer, listener);
            assertEquals(2, listener.callsToAssigned);
            assertEquals(1, listener.callsToRevoked);
        }
    }

    @ClusterTest
    public void testClassicConsumerMaxPollIntervalMsDelayInRevocation() throws InterruptedException {
        testMaxPollIntervalMsDelayInRevocation(Map.of(
            MAX_POLL_INTERVAL_MS_CONFIG, 5000,
            GROUP_PROTOCOL_CONFIG, GroupProtocol.CLASSIC.name().toLowerCase(Locale.ROOT),
            HEARTBEAT_INTERVAL_MS_CONFIG, 500,
            SESSION_TIMEOUT_MS_CONFIG, 1000,
            ENABLE_AUTO_COMMIT_CONFIG, false
        ));
    }

    @ClusterTest
    public void testAsyncConsumerMaxPollIntervalMsDelayInRevocation() throws InterruptedException {
        testMaxPollIntervalMsDelayInRevocation(Map.of(
            MAX_POLL_INTERVAL_MS_CONFIG, 5000,
            GROUP_PROTOCOL_CONFIG, GroupProtocol.CONSUMER.name().toLowerCase(Locale.ROOT),
            ENABLE_AUTO_COMMIT_CONFIG, false
        ));
    }

    private void testMaxPollIntervalMsDelayInRevocation(Map<String, Object> config) throws InterruptedException {
        var commitCompleted = new AtomicBoolean(false);
        var committedPosition = new AtomicLong(-1);

        String otherTopic = "otherTopic";
        TopicPartition tpOther = new TopicPartition(otherTopic, 0);
        sendRecords(cluster, tp, 1);
        sendRecords(cluster, tpOther, 1);

        try (Consumer<byte[], byte[]> consumer = cluster.consumer(config)) {
            var listener = new TestConsumerReassignmentListener() {
                @Override
                public void onPartitionsLost(Collection<TopicPartition> partitions) {
                    // no op
                }

                @Override
                public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
                    if (!partitions.isEmpty() && partitions.contains(tp)) {
                        // on the second rebalance (after we have joined the group initially), sleep longer
                        // than session timeout and then try a commit. We should still be in the group,
                        // so the commit should succeed
                        Utils.sleep(1500);
                        committedPosition.set(consumer.position(tp));
                        var offsets = Map.of(tp, new OffsetAndMetadata(committedPosition.get()));
                        consumer.commitSync(offsets);
                        commitCompleted.set(true);
                    }
                    super.onPartitionsRevoked(partitions);
                }
            };
            consumer.subscribe(List.of(topic), listener);

            // Consume records to ensure the rebalance completed and positions are initialized
            // (position then used in callback, triggered on next rebalance)
            awaitNonEmptyRecords(consumer, tp, 100);

            // force a rebalance to trigger an invocation of the revocation callback while in the group
            consumer.subscribe(List.of(otherTopic), listener);
            // Consume records to ensure positions for otherTopic are initialized
            // (position then used in callback, triggered on close)
            awaitNonEmptyRecords(consumer, tpOther, 100);

            assertEquals(1, committedPosition.get());
            assertTrue(commitCompleted.get());
        }
    }

    @ClusterTest
    public void testClassicConsumerMaxPollIntervalMsDelayInAssignment() throws InterruptedException {
        testMaxPollIntervalMsDelayInAssignment(Map.of(
            MAX_POLL_INTERVAL_MS_CONFIG, 5000,
            GROUP_PROTOCOL_CONFIG, GroupProtocol.CLASSIC.name().toLowerCase(Locale.ROOT),
            HEARTBEAT_INTERVAL_MS_CONFIG, 500,
            SESSION_TIMEOUT_MS_CONFIG, 1000,
            ENABLE_AUTO_COMMIT_CONFIG, false
        ));
    }

    @ClusterTest
    public void testAsyncConsumerMaxPollIntervalMsDelayInAssignment() throws InterruptedException {
        testMaxPollIntervalMsDelayInAssignment(Map.of(
            MAX_POLL_INTERVAL_MS_CONFIG, 5000,
            GROUP_PROTOCOL_CONFIG, GroupProtocol.CONSUMER.name().toLowerCase(Locale.ROOT),
            ENABLE_AUTO_COMMIT_CONFIG, false
        ));
    }

    private void testMaxPollIntervalMsDelayInAssignment(Map<String, Object> config) throws InterruptedException {
        try (Consumer<byte[], byte[]> consumer = cluster.consumer(config)) {
            var listener = new TestConsumerReassignmentListener() {
                @Override
                public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
                    // sleep longer than the session timeout, we should still be in the group after invocation
                    Utils.sleep(1500);
                    super.onPartitionsAssigned(partitions);
                }
            };
            consumer.subscribe(List.of(topic), listener);
            // rebalance to get the initial assignment
            awaitRebalance(consumer, listener);
            // We should still be in the group after this invocation
            ensureNoRebalance(consumer, listener);
        }
    }

    @ClusterTest
    public void testClassicConsumerMaxPollIntervalMsShorterThanPollTimeout() throws InterruptedException {
        testMaxPollIntervalMsShorterThanPollTimeout(Map.of(
            MAX_POLL_INTERVAL_MS_CONFIG, 1000,
            GROUP_PROTOCOL_CONFIG, GroupProtocol.CLASSIC.name().toLowerCase(Locale.ROOT),
            HEARTBEAT_INTERVAL_MS_CONFIG, 500
        ));
    }

    @ClusterTest
    public void testAsyncConsumerMaxPollIntervalMsShorterThanPollTimeout() throws InterruptedException {
        testMaxPollIntervalMsShorterThanPollTimeout(Map.of(
            MAX_POLL_INTERVAL_MS_CONFIG, 1000,
            GROUP_PROTOCOL_CONFIG, GroupProtocol.CONSUMER.name().toLowerCase(Locale.ROOT)
        ));
    }

    private void testMaxPollIntervalMsShorterThanPollTimeout(Map<String, Object> config) throws InterruptedException {
        try (Consumer<byte[], byte[]> consumer = cluster.consumer(config)) {
            var listener = new TestConsumerReassignmentListener();
            consumer.subscribe(List.of(topic), listener);

            // rebalance to get the initial assignment
            awaitRebalance(consumer, listener);
            var callsToAssignedAfterFirstRebalance = listener.callsToAssigned;

            consumer.poll(Duration.ofMillis(2000));
            // If the poll above times out, it would trigger a rebalance.
            // Leave some time for the rebalance to happen and check for the rebalance event.
            consumer.poll(Duration.ofMillis(500));
            consumer.poll(Duration.ofMillis(500));

            assertEquals(callsToAssignedAfterFirstRebalance, listener.callsToAssigned);
        }
    }

    @ClusterTest
    public void testClassicConsumerPerPartitionLeadWithMaxPollRecords() throws InterruptedException {
        testPerPartitionLeadWithMaxPollRecords(GroupProtocol.CLASSIC);
    }

    @ClusterTest
    public void testAsyncConsumerPerPartitionLeadWithMaxPollRecords() throws InterruptedException {
        testPerPartitionLeadWithMaxPollRecords(GroupProtocol.CONSUMER);
    }

    private void testPerPartitionLeadWithMaxPollRecords(GroupProtocol groupProtocol) throws InterruptedException {
        int numMessages = 1000;
        int maxPollRecords = 10;
        Map<String, Object> config = Map.of(
            GROUP_PROTOCOL_CONFIG, groupProtocol.name().toLowerCase(Locale.ROOT),
            GROUP_ID_CONFIG, "testPerPartitionLeadWithMaxPollRecords",
            CLIENT_ID_CONFIG, "testPerPartitionLeadWithMaxPollRecords",
            MAX_POLL_RECORDS_CONFIG, maxPollRecords
        );

        sendRecords(cluster, tp, numMessages);
        try (Consumer<byte[], byte[]> consumer = cluster.consumer(config)) {
            consumer.assign(List.of(tp));
            awaitNonEmptyRecords(consumer, tp, 100);

            var tags = Map.of(
                "client-id", "testPerPartitionLeadWithMaxPollRecords",
                "topic", tp.topic(),
                "partition", String.valueOf(tp.partition())
            );
            var lead = consumer.metrics()
                    .get(new MetricName("records-lead", "consumer-fetch-manager-metrics", "", tags));
            assertEquals(maxPollRecords, (Double) lead.metricValue(), "The lead should be " + maxPollRecords);
        }
    }

    @ClusterTest
    public void testClassicConsumerPerPartitionLagWithMaxPollRecords() throws InterruptedException {
        testPerPartitionLagWithMaxPollRecords(GroupProtocol.CLASSIC);
    }

    @ClusterTest
    public void testAsyncConsumerPerPartitionLagWithMaxPollRecords() throws InterruptedException {
        testPerPartitionLagWithMaxPollRecords(GroupProtocol.CONSUMER);
    }

    private void testPerPartitionLagWithMaxPollRecords(GroupProtocol groupProtocol) throws InterruptedException {
        int numMessages = 1000;
        int maxPollRecords = 10;
        Map<String, Object> config = Map.of(
            GROUP_PROTOCOL_CONFIG, groupProtocol.name().toLowerCase(Locale.ROOT),
            GROUP_ID_CONFIG, "testPerPartitionLagWithMaxPollRecords",
            CLIENT_ID_CONFIG, "testPerPartitionLagWithMaxPollRecords",
            MAX_POLL_RECORDS_CONFIG, maxPollRecords
        );
        sendRecords(cluster, tp, numMessages);
        try (Consumer<byte[], byte[]> consumer = cluster.consumer(config)) {
            consumer.assign(List.of(tp));
            var records = awaitNonEmptyRecords(consumer, tp, 100);

            var tags = Map.of(
                "client-id", "testPerPartitionLagWithMaxPollRecords",
                "topic", tp.topic(),
                "partition", String.valueOf(tp.partition())
            );
            var lag = consumer.metrics()
                    .get(new MetricName("records-lag", "consumer-fetch-manager-metrics", "", tags));

            // Count the number of records received
            var recordCount = records.count();
            assertEquals(
                numMessages - recordCount,
                (Double) lag.metricValue(),
                EPSILON,
                "The lag should be " + (numMessages - recordCount)
            );
        }
    }

    @ClusterTest
    public void runCloseClassicConsumerMultiConsumerSessionTimeoutTest() throws InterruptedException {
        runMultiConsumerSessionTimeoutTest(GroupProtocol.CLASSIC, true);
    }

    @ClusterTest
    public void runClassicConsumerMultiConsumerSessionTimeoutTest() throws InterruptedException {
        runMultiConsumerSessionTimeoutTest(GroupProtocol.CLASSIC, false);
    }

    @ClusterTest
    public void runCloseAsyncConsumerMultiConsumerSessionTimeoutTest() throws InterruptedException {
        runMultiConsumerSessionTimeoutTest(GroupProtocol.CONSUMER, true);
    }

    @ClusterTest
    public void runAsyncConsumerMultiConsumerSessionTimeoutTest() throws InterruptedException {
        runMultiConsumerSessionTimeoutTest(GroupProtocol.CONSUMER, false);
    }

    private void runMultiConsumerSessionTimeoutTest(GroupProtocol groupProtocol, boolean closeConsumer) throws InterruptedException {
        String topic1 = "topic1";
        int partitions = 6;
        Map<String, Object> config = Map.of(
            GROUP_PROTOCOL_CONFIG, groupProtocol.name().toLowerCase(Locale.ROOT),
            GROUP_ID_CONFIG, "test-group",
            MAX_POLL_INTERVAL_MS_CONFIG, 1000
        );
        // use consumers defined in this class plus one additional consumer
        // Use topic defined in this class + one additional topic
        try (Producer<byte[], byte[]> producer = cluster.producer();
             // create one more consumer and add it to the group; we will time out this consumer
             Consumer<byte[], byte[]> consumer1 = cluster.consumer(config);
             Consumer<byte[], byte[]> consumer2 = cluster.consumer(config);
             Consumer<byte[], byte[]> timeoutConsumer = cluster.consumer(config)
        ) {
            sendRecords(producer, tp, 100, System.currentTimeMillis(), -1);
            sendRecords(producer, tp2, 100, System.currentTimeMillis(), -1);

            Set<TopicPartition> subscriptions = new HashSet<>();
            subscriptions.add(tp);
            subscriptions.add(tp2);

            cluster.createTopic(topic1, partitions, (short) BROKER_COUNT);
            IntStream.range(0, partitions).forEach(partition -> {
                TopicPartition topicPartition = new TopicPartition(topic1, partition);
                sendRecords(producer, topicPartition, 100, System.currentTimeMillis(), -1);
                subscriptions.add(topicPartition);
            });

            // first subscribe consumers that are defined in this class
            List<ConsumerAssignmentPoller> consumerPollers = new ArrayList<>();
            try {
                consumerPollers.add(subscribeConsumerAndStartPolling(consumer1, List.of(topic, topic1)));
                consumerPollers.add(subscribeConsumerAndStartPolling(consumer2, List.of(topic, topic1)));

                ConsumerAssignmentPoller timeoutPoller = subscribeConsumerAndStartPolling(timeoutConsumer, List.of(topic, topic1));
                consumerPollers.add(timeoutPoller);

                // validate the initial assignment
                validateGroupAssignment(consumerPollers, subscriptions, null);

                // stop polling and close one of the consumers, should trigger partition re-assignment among alive consumers
                timeoutPoller.shutdown();
                consumerPollers.remove(timeoutPoller);
                if (closeConsumer)
                    timeoutConsumer.close();

                validateGroupAssignment(consumerPollers, subscriptions,
                        "Did not get valid assignment for partitions " + subscriptions + " after one consumer left");
            } finally {
                // done with pollers and consumers
                for (ConsumerAssignmentPoller poller : consumerPollers)
                    poller.shutdown();
            }
        }
    }

    @ClusterTest
    public void testClassicConsumerPollEventuallyReturnsRecordsWithZeroTimeout() throws InterruptedException {
        testPollEventuallyReturnsRecordsWithZeroTimeout(GroupProtocol.CLASSIC);
    }

    @ClusterTest
    public void testAsyncConsumerPollEventuallyReturnsRecordsWithZeroTimeout() throws InterruptedException {
        testPollEventuallyReturnsRecordsWithZeroTimeout(GroupProtocol.CONSUMER);
    }

    private void testPollEventuallyReturnsRecordsWithZeroTimeout(GroupProtocol groupProtocol) throws InterruptedException {
        int numMessages = 100;
        Map<String, Object> config = Map.of(GROUP_PROTOCOL_CONFIG, groupProtocol.name().toLowerCase(Locale.ROOT));
        sendRecords(cluster, tp, numMessages);

        try (Consumer<byte[], byte[]> consumer = cluster.consumer(config)) {
            consumer.subscribe(List.of(topic));
            var records = awaitNonEmptyRecords(consumer, tp, 0L);
            assertEquals(numMessages, records.count());
        }

    }

    @ClusterTest
    public void testClassicConsumerNoOffsetForPartitionExceptionOnPollZero() throws InterruptedException {
        testNoOffsetForPartitionExceptionOnPollZero(GroupProtocol.CLASSIC);
    }

    @ClusterTest
    public void testAsyncConsumerNoOffsetForPartitionExceptionOnPollZero() throws InterruptedException {
        testNoOffsetForPartitionExceptionOnPollZero(GroupProtocol.CONSUMER);
    }

    private void testNoOffsetForPartitionExceptionOnPollZero(GroupProtocol groupProtocol) throws InterruptedException {
        Map<String, Object> config = Map.of(
            GROUP_PROTOCOL_CONFIG, groupProtocol.name().toLowerCase(Locale.ROOT),
            AUTO_OFFSET_RESET_CONFIG, "none"
        );
        try (Consumer<byte[], byte[]> consumer = cluster.consumer(config)) {
            consumer.assign(List.of(tp));

            // continuous poll should eventually fail because there is no offset reset strategy set
            // (fail only when resetting positions after coordinator is known)
            waitForPollThrowException(consumer, NoOffsetForPartitionException.class);
        }
    }

    @ClusterTest
    public void testClassicConsumerRecoveryOnPollAfterDelayedRebalance() throws InterruptedException {
        testConsumerRecoveryOnPollAfterDelayedRebalance(GroupProtocol.CLASSIC);
    }

    @ClusterTest
    public void testAsyncConsumerRecoveryOnPollAfterDelayedRebalance() throws InterruptedException {
        testConsumerRecoveryOnPollAfterDelayedRebalance(GroupProtocol.CONSUMER);
    }

    public void testConsumerRecoveryOnPollAfterDelayedRebalance(GroupProtocol groupProtocol) throws InterruptedException {
        var rebalanceTimeout = 1000;
        Map<String, Object> config = Map.of(
            GROUP_PROTOCOL_CONFIG, groupProtocol.name().toLowerCase(Locale.ROOT),
            MAX_POLL_INTERVAL_MS_CONFIG, 1000,
            ENABLE_AUTO_COMMIT_CONFIG, false
        );
        try (Producer<byte[], byte[]> producer = cluster.producer();
             // Subscribe consumer that will reconcile in time on the first rebalance, but will
             // take longer than the allowed timeout in the second rebalance (onPartitionsRevoked) to get fenced by the broker.
             // The consumer should recover after being fenced (automatically rejoin the group on the next call to poll)
             Consumer<byte[], byte[]> consumer = cluster.consumer(config)
        ) {
            var numMessages = 10;
            var otherTopic = "otherTopic";
            var tpOther = new TopicPartition(otherTopic, 0);
            cluster.createTopic(otherTopic, 1, (short) BROKER_COUNT);
            sendRecords(producer, tpOther, numMessages, System.currentTimeMillis(), -1);
            sendRecords(producer, tp, numMessages, System.currentTimeMillis(), -1);

            var rebalanceTimeoutExceeded = new AtomicBoolean(false);

            var listener = new TestConsumerReassignmentListener() {
                @Override
                public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
                    if (!partitions.isEmpty() && partitions.contains(tp)) {
                        // on the second rebalance (after we have joined the group initially), sleep longer
                        // than rebalance timeout to get fenced.
                        Utils.sleep(rebalanceTimeout + 500);
                        rebalanceTimeoutExceeded.set(true);
                    }
                    super.onPartitionsRevoked(partitions);
                }
            };
            // Subscribe to get first assignment (no delays) and verify consumption
            consumer.subscribe(List.of(topic), listener);
            var records = awaitNonEmptyRecords(consumer, tp, 0L);
            assertEquals(numMessages, records.count());

            // Subscribe to different topic. This will trigger the delayed revocation exceeding rebalance timeout and get fenced
            consumer.subscribe(List.of(otherTopic), listener);
            ClientsTestUtils.pollUntilTrue(
                consumer,
                rebalanceTimeoutExceeded::get,
                "Timeout waiting for delayed callback to complete"
            );

            // Verify consumer recovers after being fenced, being able to continue consuming.
            // (The member should automatically rejoin on the next poll, with the new topic as subscription)
            records = awaitNonEmptyRecords(consumer, tpOther, 0L);
            assertEquals(numMessages, records.count());
        }
    }

    /**
     * Subscribes consumer 'consumer' to a given list of topics 'topicsToSubscribe', creates
     * consumer poller and starts polling.
     * Assumes that the consumer is not subscribed to any topics yet
     *
     * @param topicsToSubscribe topics that this consumer will subscribe to
     * @return consumer poller for the given consumer
     */
    private ConsumerAssignmentPoller subscribeConsumerAndStartPolling(
        Consumer<byte[], byte[]> consumer,
        List<String> topicsToSubscribe
    ) {
        assertEquals(0, consumer.assignment().size());
        ConsumerAssignmentPoller consumerPoller;
        consumerPoller = new ConsumerAssignmentPoller(consumer, topicsToSubscribe);
        consumerPoller.start();
        return consumerPoller;
    }

    /**
     * Check whether partition assignment is valid
     * Assumes partition assignment is valid iff
     * 1. Every consumer got assigned at least one partition
     * 2. Each partition is assigned to only one consumer
     * 3. Every partition is assigned to one of the consumers
     *
     * @param assignments set of consumer assignments; one per each consumer
     * @param partitions set of partitions that consumers subscribed to
     * @return true if partition assignment is valid
     */
    private boolean isPartitionAssignmentValid(
        List<Set<TopicPartition>> assignments,
        Set<TopicPartition> partitions
    ) {
        // check that all consumers got at least one partition
        var allNonEmptyAssignments = assignments
            .stream()
            .noneMatch(Set::isEmpty);

        if (!allNonEmptyAssignments) {
            // at least one consumer got empty assignment
            return false;
        }

        // make sure that sum of all partitions to all consumers equals total number of partitions
        var totalPartitionsInAssignments = 0;
        for (var assignment : assignments) {
            totalPartitionsInAssignments += assignment.size();
        }

        if (totalPartitionsInAssignments != partitions.size()) {
            // either same partitions got assigned to more than one consumer or some
            // partitions were not assigned
            return false;
        }

        // The above checks could miss the case where one or more partitions were assigned to more
        // than one consumer and the same number of partitions were missing from assignments.
        // Make sure that all unique assignments are the same as 'partitions'
        var uniqueAssignedPartitions = new HashSet<>();
        for (var assignment : assignments) {
            uniqueAssignedPartitions.addAll(assignment);
        }

        return uniqueAssignedPartitions.equals(partitions);
    }

    /**
     * Wait for consumers to get partition assignment and validate it.
     *
     * @param consumerPollers consumer pollers corresponding to the consumer group we are testing
     * @param subscriptions   set of all topic partitions
     * @param msg             message to print when waiting for/validating assignment fails
     */
    private void validateGroupAssignment(
        List<ConsumerAssignmentPoller> consumerPollers,
        Set<TopicPartition> subscriptions,
        String msg
    ) throws InterruptedException {
        List<Set<TopicPartition>> assignments = new ArrayList<>();
        TestUtils.waitForCondition(() -> {
            assignments.clear();
            for (ConsumerAssignmentPoller poller : consumerPollers) {
                assignments.add(poller.consumerAssignment());
            }
            return isPartitionAssignmentValid(assignments, subscriptions);
        }, GROUP_MAX_SESSION_TIMEOUT_MS * 3,
                () -> msg != null ? msg : "Did not get valid assignment for partitions " + subscriptions + ". Instead, got " + assignments
        );
    }

    private ConsumerRecords<byte[], byte[]> awaitNonEmptyRecords(
        Consumer<byte[], byte[]> consumer,
        TopicPartition partition,
        long pollTimeoutMs
    ) throws InterruptedException {
        List<ConsumerRecords<byte[], byte[]>> result = new ArrayList<>();
        TestUtils.waitForCondition(() -> {
            var records = consumer.poll(Duration.ofMillis(pollTimeoutMs));
            result.add(records);
            return !records.records(partition).isEmpty();
        }, "Consumer did not consume any messages for partition " + partition + " before timeout.");
        return result.get(result.size() - 1);
    }
}
