KAFKA-15561: Client support for new SubscriptionPattern based subscription

clients/src/main/java/org/apache/kafka/clients/consumer/Consumer.java

@@ -72,6 +72,16 @@ public interface Consumer<K, V> extends Closeable {
     */
     void subscribe(Pattern pattern);
 
+    /**
+     * @see KafkaConsumer#subscribe(SubscriptionPattern, ConsumerRebalanceListener)
+     */
+    void subscribe(SubscriptionPattern pattern, ConsumerRebalanceListener callback);
+
+    /**
+     * @see KafkaConsumer#subscribe(SubscriptionPattern)
+     */
+    void subscribe(SubscriptionPattern pattern);
+
     /**
      * @see KafkaConsumer#unsubscribe()
      */

clients/src/main/java/org/apache/kafka/clients/consumer/KafkaConsumer.java

@@ -754,6 +754,51 @@ public void subscribe(Pattern pattern) {
         delegate.subscribe(pattern);
     }
 
+    /**
+     * Subscribe to all topics matching specified pattern to get dynamically assigned partitions.
+     * The pattern matching will be done periodically against all topics existing at the time of check.
+     * This can be controlled through the {@code metadata.max.age.ms} configuration: by lowering
+     * the max metadata age, the consumer will refresh metadata more often and check for matching topics.
+     * <p>
+     * See {@link #subscribe(Collection, ConsumerRebalanceListener)} for details on the
+     * use of the {@link ConsumerRebalanceListener}. Generally rebalances are triggered when there
+     * is a change to the topics matching the provided pattern and when consumer group membership changes.
+     * Group rebalances only take place during an active call to {@link #poll(Duration)}.
+     *
+     * @param pattern SubscriptionPattern to subscribe to
+     * @param listener Non-null listener instance to get notifications on partition assignment/revocation for the
+     *                 subscribed topics
+     * @throws IllegalArgumentException If pattern or listener is null
+     * @throws IllegalStateException If {@code subscribe()} is called previously with topics, or assign is called
+     *                               previously (without a subsequent call to {@link #unsubscribe()}), or if not
+     *                               configured at-least one partition assignment strategy
+     */
+    @Override
+    public void subscribe(SubscriptionPattern pattern, ConsumerRebalanceListener callback) {
+        delegate.subscribe(pattern, callback);
+    }
+
+    /**
+     * Subscribe to all topics matching specified pattern to get dynamically assigned partitions.
+     * The pattern matching will be done on the broker against topics existing at the time of check.
+     * <p>
+     * This is a short-hand for {@link #subscribe(SubscriptionPattern, ConsumerRebalanceListener)}, which
+     * uses a no-op listener. If you need the ability to seek to particular offsets, you should prefer
+     * {@link #subscribe(SubscriptionPattern, ConsumerRebalanceListener)}, since group rebalances will cause partition offsets
+     * to be reset. You should also provide your own listener if you are doing your own offset
+     * management since the listener gives you an opportunity to commit offsets before a rebalance finishes.
+     *
+     * @param pattern SubscriptionPattern to subscribe to
+     * @throws IllegalArgumentException If pattern is null
+     * @throws IllegalStateException If {@code subscribe()} is called previously with topics, or assign is called
+     *                               previously (without a subsequent call to {@link #unsubscribe()}), or if not
+     *                               configured at-least one partition assignment strategy
+     */
+    @Override
+    public void subscribe(SubscriptionPattern pattern) {
+        delegate.subscribe(pattern);
+    }
+
     /**
      * Unsubscribe from topics currently subscribed with {@link #subscribe(Collection)} or {@link #subscribe(Pattern)}.
      * This also clears any partitions directly assigned through {@link #assign(Collection)}.

clients/src/main/java/org/apache/kafka/clients/consumer/MockConsumer.java

@@ -141,6 +141,16 @@ public synchronized void subscribe(Pattern pattern) {
         subscribe(pattern, Optional.empty());
     }
 
+    @Override
+    public void subscribe(SubscriptionPattern pattern, ConsumerRebalanceListener callback) {
+
+    }
+
+    @Override
+    public void subscribe(SubscriptionPattern pattern) {
+
+    }
+
     @Override
     public void subscribe(Collection<String> topics, final ConsumerRebalanceListener listener) {
         if (listener == null)

clients/src/main/java/org/apache/kafka/clients/consumer/SubscriptionPattern.java

@@ -0,0 +1,33 @@
+/*
+ * 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;
+
+/**
+ * A class that hold a regular expression which compatible with Google's RE2J engine. Visit
+ * <a href="https://github.com/google/re2j">this repository</a> for details on RE2J engine.
+ */
+
+public class SubscriptionPattern {
+    final private String pattern;
+    public SubscriptionPattern(final String pattern) {
+        this.pattern = pattern;
+    }
+
+    public String pattern() {
+        return this.pattern;
+    }
+}

clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java

@@ -36,6 +36,7 @@
 import org.apache.kafka.clients.consumer.OffsetAndTimestamp;
 import org.apache.kafka.clients.consumer.OffsetCommitCallback;
 import org.apache.kafka.clients.consumer.OffsetResetStrategy;
+import org.apache.kafka.clients.consumer.SubscriptionPattern;
 import org.apache.kafka.clients.consumer.internals.events.ApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.ApplicationEventHandler;
 import org.apache.kafka.clients.consumer.internals.events.ApplicationEventProcessor;
@@ -1666,6 +1667,19 @@ public void subscribe(Pattern pattern) {
         subscribeInternal(pattern, Optional.empty());
     }
 
+    @Override
+    public void subscribe(SubscriptionPattern pattern, ConsumerRebalanceListener listener) {
+        if (listener == null)
+            throw new IllegalArgumentException("RebalanceListener cannot be null");
+
+        subscribeInternal(pattern, Optional.of(listener));
+    }
+
+    @Override
+    public void subscribe(SubscriptionPattern pattern) {
+        subscribeInternal(pattern, Optional.empty());
+    }
+
     @Override
     public void subscribe(Pattern pattern, ConsumerRebalanceListener listener) {
         if (listener == null)
@@ -1730,6 +1744,21 @@ private void subscribeInternal(Pattern pattern, Optional<ConsumerRebalanceListen
         }
     }
 
+    private void subscribeInternal(SubscriptionPattern pattern, Optional<ConsumerRebalanceListener> listener) {
+        acquireAndEnsureOpen();
+        try {
+            maybeThrowInvalidGroupIdException();
+            if (pattern == null || pattern.pattern().isEmpty())
+                throw new IllegalArgumentException("Topic pattern to subscribe to cannot be " + (pattern == null ?
+                    "null" : "empty"));
+            throwIfNoAssignorsConfigured();
+            log.info("Subscribed to pattern: '{}'", pattern.pattern());
+            subscriptions.subscribe(pattern, listener);
+        } finally {
+            release();
+        }
+    }
+
     private void subscribeInternal(Collection<String> topics, Optional<ConsumerRebalanceListener> listener) {
         acquireAndEnsureOpen();
         try {

clients/src/main/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManager.java

@@ -548,7 +548,11 @@ public ConsumerGroupHeartbeatRequestData buildRequestData() {
                 }
             } else {
                 // SubscribedTopicRegex - only sent if has changed since the last heartbeat
-                //                      - not supported yet
+                String subscriptionRegex = this.subscriptions.subscriptionPattern().pattern();
+                if (!subscriptionRegex.equals(sentFields.subscribedTopicRegex)) {
+                    data.setSubscribedTopicRegex(subscriptionRegex);
+                    sentFields.subscribedTopicRegex = subscriptionRegex;
+                }
             }
 
             // ServerAssignor - only sent if has changed since the last heartbeat
@@ -591,6 +595,7 @@ static class SentFields {
             private String instanceId = null;
             private int rebalanceTimeoutMs = -1;
             private TreeSet<String> subscribedTopicNames = null;
+            private String subscribedTopicRegex = null;
             private String serverAssignor = null;
             private TreeSet<String> topicPartitions = null;
 

clients/src/main/java/org/apache/kafka/clients/consumer/internals/LegacyKafkaConsumer.java

@@ -35,6 +35,7 @@
 import org.apache.kafka.clients.consumer.OffsetAndTimestamp;
 import org.apache.kafka.clients.consumer.OffsetCommitCallback;
 import org.apache.kafka.clients.consumer.OffsetResetStrategy;
+import org.apache.kafka.clients.consumer.SubscriptionPattern;
 import org.apache.kafka.clients.consumer.internals.metrics.KafkaConsumerMetrics;
 import org.apache.kafka.common.Cluster;
 import org.apache.kafka.common.IsolationLevel;
@@ -495,6 +496,16 @@ public void subscribe(Pattern pattern) {
         subscribeInternal(pattern, Optional.empty());
     }
 
+    @Override
+    public void subscribe(SubscriptionPattern pattern, ConsumerRebalanceListener callback) {
+        log.warn("Operation not supported in new consumer group protocol");
+    }
+
+    @Override
+    public void subscribe(SubscriptionPattern pattern) {
+        log.warn("Operation not supported in new consumer group protocol");
+    }
+
     /**
      * Internal helper method for {@link #subscribe(Pattern)} and
      * {@link #subscribe(Pattern, ConsumerRebalanceListener)}

clients/src/main/java/org/apache/kafka/clients/consumer/internals/SubscriptionState.java

@@ -23,6 +23,7 @@
 import org.apache.kafka.clients.consumer.NoOffsetForPartitionException;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.clients.consumer.OffsetResetStrategy;
+import org.apache.kafka.clients.consumer.SubscriptionPattern;
 import org.apache.kafka.common.IsolationLevel;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.internals.PartitionStates;
@@ -84,6 +85,9 @@ private enum SubscriptionType {
     /* the pattern user has requested */
     private Pattern subscribedPattern;
 
+    /* RE2J compatible regex */
+    private SubscriptionPattern subscriptionPattern;
+
     /* the list of topics the user has requested */
     private Set<String> subscription;
 
@@ -108,6 +112,7 @@ public synchronized String toString() {
         return "SubscriptionState{" +
             "type=" + subscriptionType +
             ", subscribedPattern=" + subscribedPattern +
+            ", subscriptionPattern=" + subscriptionPattern +
             ", subscription=" + String.join(",", subscription) +
             ", groupSubscription=" + String.join(",", groupSubscription) +
             ", defaultResetStrategy=" + defaultResetStrategy +
@@ -136,6 +141,7 @@ public SubscriptionState(LogContext logContext, OffsetResetStrategy defaultReset
         this.assignment = new PartitionStates<>();
         this.groupSubscription = new HashSet<>();
         this.subscribedPattern = null;
+        this.subscriptionPattern = null;
         this.subscriptionType = SubscriptionType.NONE;
     }
 
@@ -174,6 +180,12 @@ public synchronized void subscribe(Pattern pattern, Optional<ConsumerRebalanceLi
         this.subscribedPattern = pattern;
     }
 
+    public synchronized void subscribe(SubscriptionPattern pattern, Optional<ConsumerRebalanceListener> listener) {
+        registerRebalanceListener(listener);
+        setSubscriptionType(SubscriptionType.AUTO_PATTERN);
+        this.subscriptionPattern = pattern;
+    }
+
     public synchronized boolean subscribeFromPattern(Set<String> topics) {
         if (subscriptionType != SubscriptionType.AUTO_PATTERN)
             throw new IllegalArgumentException("Attempt to subscribe from pattern while subscription type set to " +
@@ -306,6 +318,7 @@ public synchronized void unsubscribe() {
         this.groupSubscription = Collections.emptySet();
         this.assignment.clear();
         this.subscribedPattern = null;
+        this.subscriptionPattern = null;
         this.subscriptionType = SubscriptionType.NONE;
         this.assignmentId++;
     }
@@ -328,6 +341,10 @@ public synchronized Set<String> subscription() {
         return Collections.emptySet();
     }
 
+    public synchronized SubscriptionPattern subscriptionPattern() {
+        return this.subscriptionPattern;
+    }
+
     public synchronized Set<TopicPartition> pausedPartitions() {
         return collectPartitions(TopicPartitionState::isPaused);
     }

clients/src/main/resources/common/message/ConsumerGroupHeartbeatRequest.json

@@ -35,6 +35,8 @@
       "about": "-1 if it didn't change since the last heartbeat; the maximum time in milliseconds that the coordinator will wait on the member to revoke its partitions otherwise." },
     { "name": "SubscribedTopicNames", "type": "[]string", "versions": "0+", "nullableVersions": "0+", "default": "null", "entityType": "topicName",
       "about": "null if it didn't change since the last heartbeat; the subscribed topic names otherwise." },
+    { "name": "SubscribedTopicRegex", "type": "string", "versions": "0+", "nullableVersions": "0+", "default": "null",
+      "about": "null if it didn't change since the last heartbeat; the subscribed topic regex otherwise" },
     { "name": "ServerAssignor", "type": "string", "versions": "0+", "nullableVersions": "0+", "default": "null",
       "about": "null if not used or if it didn't change since the last heartbeat; the server side assignor to use otherwise." },
     { "name": "TopicPartitions", "type": "[]TopicPartitions", "versions": "0+", "nullableVersions": "0+", "default": "null",

clients/src/test/java/org/apache/kafka/clients/consumer/KafkaConsumerTest.java

@@ -485,6 +485,13 @@ public void testSubscriptionOnNullPattern(GroupProtocol groupProtocol) {
         assertThrows(IllegalArgumentException.class,
             () -> consumer.subscribe((Pattern) null));
     }
+    @ParameterizedTest
+    @EnumSource(value = GroupProtocol.class, names = "CONSUMER")
+    public void testSubscriptionOnNullSubscriptionPattern(GroupProtocol groupProtocol) {
+        consumer = newConsumer(groupProtocol, groupId);
+        assertThrows(IllegalArgumentException.class,
+            () -> consumer.subscribe((SubscriptionPattern) null));
+    }
 
     @ParameterizedTest
     @EnumSource(GroupProtocol.class)
@@ -494,6 +501,14 @@ public void testSubscriptionOnEmptyPattern(GroupProtocol groupProtocol) {
             () -> consumer.subscribe(Pattern.compile("")));
     }
 
+    @ParameterizedTest
+    @EnumSource(value = GroupProtocol.class, names = "CONSUMER")
+    public void testSubscriptionOnEmptySubscriptionPattern(GroupProtocol groupProtocol) {
+        consumer = newConsumer(groupProtocol, groupId);
+        assertThrows(IllegalArgumentException.class,
+                () -> consumer.subscribe(new SubscriptionPattern("")));
+    }
+
     @ParameterizedTest
     @EnumSource(GroupProtocol.class)
     public void testSubscriptionWithEmptyPartitionAssignment(GroupProtocol groupProtocol) {

clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java

@@ -30,6 +30,7 @@
 import org.apache.kafka.clients.consumer.OffsetResetStrategy;
 import org.apache.kafka.clients.consumer.RetriableCommitFailedException;
 import org.apache.kafka.clients.consumer.RoundRobinAssignor;
+import org.apache.kafka.clients.consumer.SubscriptionPattern;
 import org.apache.kafka.clients.consumer.internals.events.ApplicationEventHandler;
 import org.apache.kafka.clients.consumer.internals.events.AssignmentChangeApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEvent;
@@ -582,6 +583,23 @@ public void testPartitionRevocationOnClose() {
         assertTrue(subscriptions.assignedPartitions().isEmpty());
         assertEquals(1, listener.revokedCount);
     }
+    @Test
+    public void testSubscribeToSubscriptionPattern() {
+        MockRebalanceListener listener = new MockRebalanceListener();
+        SubscriptionState subscriptions = new SubscriptionState(new LogContext(), OffsetResetStrategy.NONE);
+        consumer = newConsumer(
+                mock(FetchBuffer.class),
+                mock(ConsumerInterceptors.class),
+                mock(ConsumerRebalanceListenerInvoker.class),
+                subscriptions,
+                singletonList(new RoundRobinAssignor()),
+                "group-id",
+                "client-id");
+
+        SubscriptionPattern pattern = new SubscriptionPattern("t*");
+        consumer.subscribe(pattern, listener);
+        assertEquals(subscriptions.subscriptionPattern(), pattern);
+    }
 
     @Test
     public void testFailedPartitionRevocationOnClose() {

clients/src/test/java/org/apache/kafka/clients/consumer/internals/HeartbeatRequestManagerTest.java

@@ -319,6 +319,7 @@ public void testValidateConsumerGroupHeartbeatRequest(final short version) {
         assertEquals(memberEpoch, heartbeatRequest.data().memberEpoch());
         assertEquals(DEFAULT_MAX_POLL_INTERVAL_MS, heartbeatRequest.data().rebalanceTimeoutMs());
         assertEquals(subscribedTopics, heartbeatRequest.data().subscribedTopicNames());
+        assertEquals(null, heartbeatRequest.data().subscribedTopicRegex());
         assertEquals(DEFAULT_GROUP_INSTANCE_ID, heartbeatRequest.data().instanceId());
         assertEquals(DEFAULT_REMOTE_ASSIGNOR, heartbeatRequest.data().serverAssignor());
     }
@@ -477,6 +478,7 @@ public void testHeartbeatState() {
         assertNull(data.instanceId());
         assertEquals(ConsumerTestBuilder.DEFAULT_MAX_POLL_INTERVAL_MS, data.rebalanceTimeoutMs());
         assertEquals(Collections.emptyList(), data.subscribedTopicNames());
+        assertEquals(null, data.subscribedTopicRegex());
         assertEquals(ConsumerTestBuilder.DEFAULT_REMOTE_ASSIGNOR, data.serverAssignor());
         assertEquals(Collections.emptyList(), data.topicPartitions());
         membershipManager.onHeartbeatRequestSent();
@@ -491,6 +493,7 @@ public void testHeartbeatState() {
         assertNull(data.instanceId());
         assertEquals(-1, data.rebalanceTimeoutMs());
         assertNull(data.subscribedTopicNames());
+        assertNull(data.subscribedTopicRegex());
         assertNull(data.serverAssignor());
         assertNull(data.topicPartitions());
         membershipManager.onHeartbeatRequestSent();
@@ -508,6 +511,7 @@ public void testHeartbeatState() {
         assertNull(data.instanceId());
         assertEquals(-1, data.rebalanceTimeoutMs());
         assertEquals(Collections.singletonList(topic), data.subscribedTopicNames());
+        assertNull(data.subscribedTopicRegex());
         assertNull(data.serverAssignor());
         assertNull(data.topicPartitions());
         membershipManager.onHeartbeatRequestSent();

clients/src/test/java/org/apache/kafka/clients/consumer/internals/OffsetsRequestManagerTest.java

@@ -86,7 +86,6 @@
 import static org.mockito.Mockito.when;
 
 public class OffsetsRequestManagerTest {
-
     private OffsetsRequestManager requestManager;
     private ConsumerMetadata metadata;
     private SubscriptionState subscriptionState;