[KIP-848] Tests for: dev_kip848_fix_fast_subscribe_or_unsubscribe (confluentinc#5026)

Author: emasab

src/rd.h

@@ -227,6 +227,12 @@ static RD_INLINE RD_UNUSED char *rd_strndup(const char *s, size_t len) {
 #define RD_ARRAYSIZE(A)           RD_ARRAY_SIZE(A)
 #define RD_SIZEOF(TYPE, MEMBER)   sizeof(((TYPE *)NULL)->MEMBER)
 #define RD_OFFSETOF(TYPE, MEMBER) ((size_t) & (((TYPE *)NULL)->MEMBER))
+/** Array foreach */
+#define RD_ARRAY_FOREACH_INDEX(ELEM, ARRAY, INDEX)                             \
+        for ((INDEX = 0, (ELEM) = (ARRAY)[INDEX]);                             \
+             INDEX < RD_ARRAY_SIZE(ARRAY);                                     \
+             (ELEM) =                                                          \
+                 (++INDEX < RD_ARRAY_SIZE(ARRAY) ? (ARRAY)[INDEX] : (ELEM)))
 
 /**
  * Returns the 'I'th array element from static sized array 'A'

tests/0022-consume_batch.c

@@ -259,7 +259,9 @@ static void do_test_consume_batch_non_existent_topic(void) {
 
 int main_0022_consume_batch(int argc, char **argv) {
         do_test_consume_batch();
-        /* FIXME: this must be implemented in KIP-848 for compatibility. */
+        /* Subscribing to a non-existent topic doesn't generate a
+         * UNKNOWN_TOPIC_OR_PART error with `consumer` group rebalance protocol.
+         */
         if (test_consumer_group_protocol_classic()) {
                 do_test_consume_batch_non_existent_topic();
         }

tests/0045-subscribe_update.c

@@ -821,6 +821,108 @@ static void do_test_resubscribe_with_regex() {
         SUB_TEST_PASS();
 }
 
+/**
+ * @brief Create many topics and apply several subscription
+ *        updates, unsubscribing and re-subscribing too.
+ *        After changing some subscriptions verifies that the assignment
+ *        corresponds to last one.
+ *
+ * @param with_rebalance_cb Use a rebalance callback to perform the assignment.
+ *                          It needs to poll the consumer when awaiting for the
+ *                          assignment in this case.
+ */
+static void do_test_subscribe_many_updates(rd_bool_t with_rebalance_cb) {
+#define TOPIC_CNT 100
+        char *topics[TOPIC_CNT] = {0};
+        char *topic;
+        char *group;
+        size_t i;
+        rd_kafka_t *rk;
+        rd_kafka_conf_t *conf;
+        const int partition_cnt = 4;
+
+        SUB_TEST("%s", with_rebalance_cb ? "with rebalance callback"
+                                         : "without rebalance callback");
+
+        RD_ARRAY_FOREACH_INDEX(topic, topics, i) {
+                char topic_i[17];
+                rd_snprintf(topic_i, sizeof(topic_i), "topic%" PRIusz, i);
+                topics[i] = rd_strdup(test_mk_topic_name(topic_i, 1));
+        };
+        group = topics[0];
+
+        test_conf_init(&conf, NULL, 60);
+        if (with_rebalance_cb)
+                rd_kafka_conf_set_rebalance_cb(conf, test_rebalance_cb);
+        rk = test_create_consumer(group, NULL, conf, NULL);
+
+        TEST_SAY("Creating %d topics\n", TOPIC_CNT);
+        TEST_CALL_ERR__(test_CreateTopics_simple(rk, NULL, topics, TOPIC_CNT,
+                                                 partition_cnt, NULL));
+        test_wait_topic_exists(rk, topics[TOPIC_CNT - 1], 5000);
+        /* Give the cluster some more time to propagate metadata
+         * for TOPICS_CNT topics */
+        rd_sleep(1);
+
+        RD_ARRAY_FOREACH_INDEX(topic, topics, i) {
+                const int max_subscription_size = 5;
+                size_t j;
+                int k;
+                int subscription_size =
+                    RD_MIN(max_subscription_size, TOPIC_CNT - i);
+                int expected_assignment_cnt = subscription_size * partition_cnt;
+                rd_kafka_topic_partition_list_t *expected_assignment = NULL;
+
+                rd_kafka_topic_partition_list_t *subscription =
+                    rd_kafka_topic_partition_list_new(subscription_size);
+                rd_bool_t check_expected_assignment =
+                    (i % 5 == 0 || i == TOPIC_CNT - 1);
+                rd_bool_t do_unsubscribe = i % 7 == 0;
+
+                if (check_expected_assignment)
+                        expected_assignment = rd_kafka_topic_partition_list_new(
+                            expected_assignment_cnt);
+
+                for (j = i; j < i + subscription_size; j++) {
+                        rd_kafka_topic_partition_list_add(
+                            subscription, topics[j], RD_KAFKA_PARTITION_UA);
+                        /* We unsubscribe every 7 iteration and
+                         * we check assignments every 5 iteration so
+                         * at 7 * 5 we unsubscribe and check that assignment
+                         * is empty. */
+                        if (check_expected_assignment && !do_unsubscribe)
+                                for (k = 0; k < partition_cnt; k++)
+                                        rd_kafka_topic_partition_list_add(
+                                            expected_assignment, topics[j], k);
+                }
+                TEST_CALL_ERR__(rd_kafka_subscribe(rk, subscription));
+                rd_kafka_topic_partition_list_destroy(subscription);
+
+                if (do_unsubscribe)
+                        TEST_CALL_ERR__(rd_kafka_unsubscribe(rk));
+                if (check_expected_assignment) {
+                        test_consumer_wait_assignment_topic_partition_list(
+                            rk,
+                            /* poll when we have a rebalance callback */
+                            with_rebalance_cb, expected_assignment, 30000);
+                        rd_kafka_topic_partition_list_destroy(
+                            expected_assignment);
+                }
+        };
+
+        TEST_CALL_ERR__(
+            test_DeleteTopics_simple(rk, NULL, topics, TOPIC_CNT, NULL));
+        RD_ARRAY_FOREACH_INDEX(topic, topics, i) {
+                rd_free(topic);
+        };
+
+        test_consumer_close(rk);
+        rd_kafka_destroy(rk);
+
+        SUB_TEST_PASS();
+#undef TOPIC_CNT
+}
+
 int main_0045_subscribe_update(int argc, char **argv) {
 
         if (!test_can_create_topics(1))
@@ -864,6 +966,12 @@ int main_0045_resubscribe_with_regex(int argc, char **argv) {
         return 0;
 }
 
+int main_0045_subscribe_many_updates(int argc, char **argv) {
+        do_test_subscribe_many_updates(rd_false);
+        do_test_subscribe_many_updates(rd_true);
+        return 0;
+}
+
 int main_0045_subscribe_update_racks_mock(int argc, char **argv) {
         int use_replica_rack = 0;
         int use_client_rack  = 0;

tests/test.c

@@ -158,6 +158,7 @@ _TEST_DECL(0045_subscribe_update_non_exist_and_partchange);
 _TEST_DECL(0045_subscribe_update_mock);
 _TEST_DECL(0045_subscribe_update_racks_mock);
 _TEST_DECL(0045_resubscribe_with_regex);
+_TEST_DECL(0045_subscribe_many_updates);
 _TEST_DECL(0046_rkt_cache);
 _TEST_DECL(0047_partial_buf_tmout);
 _TEST_DECL(0048_partitioner);
@@ -388,6 +389,7 @@ struct test tests[] = {
     _TEST(0045_subscribe_update_mock, TEST_F_LOCAL),
     _TEST(0045_subscribe_update_racks_mock, TEST_F_LOCAL),
     _TEST(0045_resubscribe_with_regex, 0, TEST_BRKVER(0, 9, 0, 0)),
+    _TEST(0045_subscribe_many_updates, 0, TEST_BRKVER(0, 10, 2, 0)),
     _TEST(0046_rkt_cache, TEST_F_LOCAL),
     _TEST(0047_partial_buf_tmout, TEST_F_KNOWN_ISSUE),
     _TEST(0048_partitioner,
@@ -3114,6 +3116,106 @@ void test_consumer_verify_assignment0(const char *func,
         rd_kafka_topic_partition_list_destroy(assignment);
 }
 
+/**
+ * @brief Verify that the consumer's assignment matches the expected assignment.
+ *        passed as a topic partition list in \p expected_assignment .
+ */
+rd_bool_t test_consumer_verify_assignment_topic_partition_list0(
+    const char *func,
+    int line,
+    rd_kafka_t *rk,
+    const rd_kafka_topic_partition_list_t *expected_assignment) {
+        rd_kafka_topic_partition_list_t *assignment,
+            *expected_assignment_copy = NULL;
+        rd_kafka_resp_err_t err;
+        int i;
+        rd_bool_t ret = rd_true;
+
+        if ((err = rd_kafka_assignment(rk, &assignment)))
+                TEST_FAIL("%s:%d: Failed to get assignment for %s: %s", func,
+                          line, rd_kafka_name(rk), rd_kafka_err2str(err));
+
+        TEST_SAYL(4, "%s assignment (%d partition(s)):\n", rd_kafka_name(rk),
+                  assignment->cnt);
+        for (i = 0; i < assignment->cnt; i++)
+                TEST_SAYL(4, " %s [%" PRId32 "]\n", assignment->elems[i].topic,
+                          assignment->elems[i].partition);
+
+        if (assignment->cnt != expected_assignment->cnt) {
+                ret = rd_false;
+                goto done;
+        }
+
+        expected_assignment_copy =
+            rd_kafka_topic_partition_list_copy(expected_assignment);
+        rd_kafka_topic_partition_list_sort(assignment, NULL, NULL);
+        rd_kafka_topic_partition_list_sort(expected_assignment_copy, NULL,
+                                           NULL);
+
+        for (i = 0; i < assignment->cnt; i++) {
+                if (strcmp(assignment->elems[i].topic,
+                           expected_assignment_copy->elems[i].topic) ||
+                    assignment->elems[i].partition !=
+                        expected_assignment_copy->elems[i].partition) {
+                        ret = rd_false;
+                        goto done;
+                }
+        }
+
+done:
+        RD_IF_FREE(expected_assignment_copy,
+                   rd_kafka_topic_partition_list_destroy);
+        rd_kafka_topic_partition_list_destroy(assignment);
+        return ret;
+}
+
+/**
+ * @brief Wait until the consumer's assignment matches the expected assignment.
+ *        passed as a topic partition list in \p expected_assignment .
+ *        Polling if \p do_poll is true, otherwise sleeps.
+ *        Until \p timeout_ms milliseconds.
+ */
+void test_consumer_wait_assignment_topic_partition_list0(
+    const char *func,
+    int line,
+    rd_kafka_t *rk,
+    rd_bool_t do_poll,
+    const rd_kafka_topic_partition_list_t *expected_assignment,
+    int timeout_ms) {
+        int i;
+        rd_ts_t end        = test_clock() + timeout_ms * 1000;
+        rd_bool_t verified = rd_false;
+
+        TEST_SAY("Verifying assignment\n");
+        TEST_SAYL(4, "%s expected assignment (%d partition(s)):\n",
+                  rd_kafka_name(rk), expected_assignment->cnt);
+        for (i = 0; i < expected_assignment->cnt; i++)
+                TEST_SAYL(4, " %s [%" PRId32 "]\n",
+                          expected_assignment->elems[i].topic,
+                          expected_assignment->elems[i].partition);
+
+        do {
+                verified =
+                    test_consumer_verify_assignment_topic_partition_list0(
+                        func, line, rk, expected_assignment);
+                if (verified)
+                        break;
+
+                if (do_poll)
+                        test_consumer_poll_once(rk, NULL, 100);
+                else
+                        rd_usleep(100 * 1000, NULL);
+
+        } while (test_clock() < end);
+
+        if (!verified) {
+                TEST_FAIL(
+                    "%s:%d: Expected assignment not found in %s's "
+                    "assignment within timeout %d",
+                    func, line, rd_kafka_name(rk), timeout_ms);
+        }
+        TEST_SAY("Verified assignment\n");
+}
 
 
 /**

tests/test.h

@@ -676,6 +676,28 @@ void test_consumer_verify_assignment0(const char *func,
         test_consumer_verify_assignment0(__FUNCTION__, __LINE__, rk,           \
                                          fail_immediately, __VA_ARGS__)
 
+rd_bool_t test_consumer_verify_assignment_topic_partition_list0(
+    const char *func,
+    int line,
+    rd_kafka_t *rk,
+    const rd_kafka_topic_partition_list_t *expected_assignment);
+#define test_consumer_verify_assignment_topic_partition_list(                  \
+    rk, expected_assignment)                                                   \
+        test_consumer_verify_assignment_topic_partition_list0(                 \
+            __FUNCTION__, __LINE__, rk, expected_assignment)
+
+void test_consumer_wait_assignment_topic_partition_list0(
+    const char *func,
+    int line,
+    rd_kafka_t *rk,
+    rd_bool_t do_poll,
+    const rd_kafka_topic_partition_list_t *expected_assignment,
+    int timeout_ms);
+#define test_consumer_wait_assignment_topic_partition_list(                    \
+    rk, do_poll, expected_assignment, timeout_ms)                              \
+        test_consumer_wait_assignment_topic_partition_list0(                   \
+            __FUNCTION__, __LINE__, rk, do_poll, expected_assignment,          \
+            timeout_ms)
 
 void test_consumer_assign(const char *what,
                           rd_kafka_t *rk,