KAFKA-14950: implement assign() and assignment()

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

@@ -121,7 +121,7 @@ public DefaultBackgroundThread(final Time time,
             this.errorEventHandler = new ErrorEventHandler(this.backgroundEventQueue);
             this.groupState = new GroupState(rebalanceConfig);
             this.requestManagerRegistry = Collections.unmodifiableMap(buildRequestManagerRegistry(logContext));
-            this.applicationEventProcessor = new ApplicationEventProcessor(backgroundEventQueue, requestManagerRegistry);
+            this.applicationEventProcessor = new ApplicationEventProcessor(backgroundEventQueue, requestManagerRegistry, metadata);
         } catch (final Exception e) {
             close();
             throw new KafkaException("Failed to construct background processor", e.getCause());

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

@@ -32,7 +32,9 @@
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEvent;
 import org.apache.kafka.clients.consumer.internals.events.CommitApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.EventHandler;
+import org.apache.kafka.clients.consumer.internals.events.MetadataUpdateApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.OffsetFetchApplicationEvent;
+import org.apache.kafka.clients.consumer.internals.events.UnsubscribeApplicationEvent;
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.Metric;
 import org.apache.kafka.common.MetricName;
@@ -58,6 +60,7 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
@@ -497,7 +500,7 @@ public void commitSync(Map<TopicPartition, OffsetAndMetadata> offsets, Duration
 
     @Override
     public Set<TopicPartition> assignment() {
-        throw new KafkaException("method not implemented");
+        return Collections.unmodifiableSet(this.subscriptions.assignedPartitions());
     }
 
     /**
@@ -522,7 +525,35 @@ public void subscribe(Collection<String> topics, ConsumerRebalanceListener callb
 
     @Override
     public void assign(Collection<TopicPartition> partitions) {
-        throw new KafkaException("method not implemented");
+        if (partitions == null) {
+            throw new IllegalArgumentException("Topic partitions collection to assign to cannot be null");
+        }
+
+        if (partitions.isEmpty()) {
+            this.unsubscribe();
+            return;
+        }
+
+        for (TopicPartition tp : partitions) {
+            String topic = (tp != null) ? tp.topic() : null;
+            if (Utils.isBlank(topic))
+                throw new IllegalArgumentException("Topic partitions to assign to cannot have null or empty topic");
+        }
+        // TODO: implement fetcher
+        // fetcher.clearBufferedDataForUnassignedPartitions(partitions);
+
+        // make sure the offsets of topic partitions the consumer is unsubscribing from
+        // are committed since there will be no following rebalance
+        commit(subscriptions.allConsumed());
+
+        log.info("Assigned to partition(s): {}", Utils.join(partitions, ", "));
+        if (this.subscriptions.assignFromUser(new HashSet<>(partitions)))
+           updateMetadata(time.milliseconds());
+    }
+
+    private void updateMetadata(long milliseconds) {
+        final MetadataUpdateApplicationEvent event = new MetadataUpdateApplicationEvent(milliseconds);
+        eventHandler.add(event);
     }
 
     @Override
@@ -537,7 +568,9 @@ public void subscribe(Pattern pattern) {
 
     @Override
     public void unsubscribe() {
-        throw new KafkaException("method not implemented");
+        // fetcher.clearBufferedDataForUnassignedPartitions(Collections.emptySet());
+        eventHandler.add(new UnsubscribeApplicationEvent());
+        this.subscriptions.unsubscribe();
     }
 
     @Override

clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEvent.java

@@ -36,6 +36,6 @@ public String toString() {
         return type + " ApplicationEvent";
     }
     public enum Type {
-        NOOP, COMMIT, POLL, FETCH_COMMITTED_OFFSET,
+        NOOP, COMMIT, POLL, FETCH_COMMITTED_OFFSET, METADATA_UPDATE, UNSUBSCRIBE,
     }
 }

clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessor.java

@@ -17,6 +17,7 @@
 package org.apache.kafka.clients.consumer.internals.events;
 
 import org.apache.kafka.clients.consumer.internals.CommitRequestManager;
+import org.apache.kafka.clients.consumer.internals.ConsumerMetadata;
 import org.apache.kafka.clients.consumer.internals.NoopBackgroundEvent;
 import org.apache.kafka.clients.consumer.internals.RequestManager;
 import org.apache.kafka.common.KafkaException;
@@ -29,12 +30,15 @@
 public class ApplicationEventProcessor {
     private final BlockingQueue<BackgroundEvent> backgroundEventQueue;
     private final Map<RequestManager.Type, Optional<RequestManager>> registry;
+    private final ConsumerMetadata metadata;
 
     public ApplicationEventProcessor(
-            final BlockingQueue<BackgroundEvent> backgroundEventQueue,
-            final Map<RequestManager.Type, Optional<RequestManager>> requestManagerRegistry) {
+        final BlockingQueue<BackgroundEvent> backgroundEventQueue,
+        final Map<RequestManager.Type, Optional<RequestManager>> requestManagerRegistry,
+        final ConsumerMetadata metadata) {
         this.backgroundEventQueue = backgroundEventQueue;
         this.registry = requestManagerRegistry;
+        this.metadata = metadata;
     }
 
     public boolean process(final ApplicationEvent event) {
@@ -48,6 +52,10 @@ public boolean process(final ApplicationEvent event) {
                 return process((PollApplicationEvent) event);
             case FETCH_COMMITTED_OFFSET:
                 return process((OffsetFetchApplicationEvent) event);
+            case METADATA_UPDATE:
+                return process((MetadataUpdateApplicationEvent) event);
+            case UNSUBSCRIBE:
+                return process((UnsubscribeApplicationEvent) event);
         }
         return false;
     }
@@ -106,4 +114,17 @@ private boolean process(final OffsetFetchApplicationEvent event) {
         manager.addOffsetFetchRequest(event.partitions);
         return true;
     }
+
+    private boolean process(final MetadataUpdateApplicationEvent event) {
+        metadata.requestUpdateForNewTopics();
+        return true;
+    }
+
+    private boolean process(final UnsubscribeApplicationEvent event) {
+        /*
+                this.coordinator.onLeavePrepare();
+                this.coordinator.maybeLeaveGroup("the consumer unsubscribed from all topics");
+         */
+        return true;
+    }
 }

clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/MetadataUpdateApplicationEvent.java

@@ -0,0 +1,27 @@
+/*
+ * 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.internals.events;
+
+public class MetadataUpdateApplicationEvent extends ApplicationEvent {
+
+    private final long timestamp;
+
+    public MetadataUpdateApplicationEvent(final long timestamp) {
+        super(Type.METADATA_UPDATE);
+        this.timestamp = timestamp;
+    }
+}

clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/UnsubscribeApplicationEvent.java

@@ -0,0 +1,24 @@
+/*
+ * 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.internals.events;
+
+public class UnsubscribeApplicationEvent extends ApplicationEvent {
+
+    public UnsubscribeApplicationEvent() {
+        super(Type.UNSUBSCRIBE);
+    }
+}

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

@@ -21,6 +21,8 @@
 import org.apache.kafka.clients.consumer.internals.events.ApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.ApplicationEventProcessor;
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEvent;
+import org.apache.kafka.clients.consumer.internals.events.CommitApplicationEvent;
+import org.apache.kafka.clients.consumer.internals.events.MetadataUpdateApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.NoopApplicationEvent;
 import org.apache.kafka.common.message.FindCoordinatorRequestData;
 import org.apache.kafka.common.requests.FindCoordinatorRequest;
@@ -34,6 +36,8 @@
 
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Properties;
 import java.util.concurrent.BlockingQueue;
@@ -44,6 +48,7 @@
 import static org.apache.kafka.clients.consumer.ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
@@ -58,7 +63,7 @@ public class DefaultBackgroundThreadTest {
     private NetworkClientDelegate networkClient;
     private BlockingQueue<BackgroundEvent> backgroundEventsQueue;
     private BlockingQueue<ApplicationEvent> applicationEventsQueue;
-    private ApplicationEventProcessor processor;
+    private ApplicationEventProcessor applicationEventProcessor;
     private CoordinatorRequestManager coordinatorManager;
     private ErrorEventHandler errorEventHandler;
     private int requestTimeoutMs = 500;
@@ -73,7 +78,7 @@ public void setup() {
         this.networkClient = mock(NetworkClientDelegate.class);
         this.applicationEventsQueue = (BlockingQueue<ApplicationEvent>) mock(BlockingQueue.class);
         this.backgroundEventsQueue = (BlockingQueue<BackgroundEvent>) mock(BlockingQueue.class);
-        this.processor = mock(ApplicationEventProcessor.class);
+        this.applicationEventProcessor = mock(ApplicationEventProcessor.class);
         this.coordinatorManager = mock(CoordinatorRequestManager.class);
         this.errorEventHandler = mock(ErrorEventHandler.class);
         GroupRebalanceConfig rebalanceConfig = new GroupRebalanceConfig(
@@ -106,7 +111,37 @@ public void testApplicationEvent() {
         ApplicationEvent e = new NoopApplicationEvent("noop event");
         this.applicationEventsQueue.add(e);
         backgroundThread.runOnce();
-        verify(processor, times(1)).process(e);
+        verify(applicationEventProcessor, times(1)).process(e);
+        backgroundThread.close();
+    }
+
+    @Test
+    public void testMetadataUpdateEvent() {
+        this.applicationEventsQueue = new LinkedBlockingQueue<>();
+        this.backgroundEventsQueue = new LinkedBlockingQueue<>();
+        this.applicationEventProcessor = new ApplicationEventProcessor(this.backgroundEventsQueue, mockRequestManagerRegistry(),
+            metadata);
+        when(coordinatorManager.poll(anyLong())).thenReturn(mockPollCoordinatorResult());
+        when(commitManager.poll(anyLong())).thenReturn(mockPollCommitResult());
+        DefaultBackgroundThread backgroundThread = mockBackgroundThread();
+        ApplicationEvent e = new MetadataUpdateApplicationEvent(time.milliseconds());
+        this.applicationEventsQueue.add(e);
+        backgroundThread.runOnce();
+        verify(metadata).requestUpdateForNewTopics();
+        backgroundThread.close();
+    }
+
+    @Test
+    public void testCommitEvent() {
+        this.applicationEventsQueue = new LinkedBlockingQueue<>();
+        this.backgroundEventsQueue = new LinkedBlockingQueue<>();
+        when(coordinatorManager.poll(anyLong())).thenReturn(mockPollCoordinatorResult());
+        when(commitManager.poll(anyLong())).thenReturn(mockPollCommitResult());
+        DefaultBackgroundThread backgroundThread = mockBackgroundThread();
+        ApplicationEvent e = new CommitApplicationEvent(new HashMap<>());
+        this.applicationEventsQueue.add(e);
+        backgroundThread.runOnce();
+        verify(applicationEventProcessor).process(any(CommitApplicationEvent.class));
         backgroundThread.close();
     }
 
@@ -136,6 +171,13 @@ void testPollResultTimer() {
         assertEquals(10, backgroundThread.handlePollResult(failure));
     }
 
+    private Map<RequestManager.Type, Optional<RequestManager>> mockRequestManagerRegistry() {
+        Map<RequestManager.Type, Optional<RequestManager>> registry = new HashMap<>();
+        registry.put(RequestManager.Type.COORDINATOR, Optional.of(coordinatorManager));
+        registry.put(RequestManager.Type.COMMIT, Optional.of(coordinatorManager));
+        return registry;
+    }
+
     private static NetworkClientDelegate.UnsentRequest findCoordinatorUnsentRequest(
             final Time time,
             final long timeout
@@ -162,7 +204,7 @@ private DefaultBackgroundThread mockBackgroundThread() {
                 applicationEventsQueue,
                 backgroundEventsQueue,
                 this.errorEventHandler,
-                processor,
+            applicationEventProcessor,
                 this.metadata,
                 this.networkClient,
                 this.groupState,

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

@@ -20,9 +20,11 @@
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.clients.consumer.OffsetCommitCallback;
+import org.apache.kafka.clients.consumer.OffsetResetStrategy;
+import org.apache.kafka.clients.consumer.internals.events.CommitApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.EventHandler;
+import org.apache.kafka.clients.consumer.internals.events.MetadataUpdateApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.OffsetFetchApplicationEvent;
-import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.InvalidGroupIdException;
 import org.apache.kafka.common.internals.ClusterResourceListeners;
@@ -39,20 +41,23 @@
 import org.mockito.ArgumentMatchers;
 
 import java.time.Duration;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 
+import static java.util.Collections.singleton;
 import static org.apache.kafka.clients.consumer.ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG;
 import static org.apache.kafka.clients.consumer.ConsumerConfig.DEFAULT_API_TIMEOUT_MS_CONFIG;
 import static org.apache.kafka.clients.consumer.ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG;
 import static org.apache.kafka.clients.consumer.ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG;
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
@@ -157,9 +162,41 @@ public void testCommitted() {
     }
 
     @Test
-    public void testUnimplementedException() {
+    public void testAssign() {
+        this.subscriptions = new SubscriptionState(logContext, OffsetResetStrategy.EARLIEST);
+        this.consumer = newConsumer(time, new StringDeserializer(), new StringDeserializer());
+        final TopicPartition tp = new TopicPartition("foo", 3);
+        consumer.assign(singleton(tp));
+        assertTrue(consumer.subscription().isEmpty());
+        assertTrue(consumer.assignment().contains(tp));
+        verify(eventHandler).add(any(CommitApplicationEvent.class));
+        verify(eventHandler).add(any(MetadataUpdateApplicationEvent.class));
+    }
+
+    @Test
+    public void testAssignOnNullTopicPartition() {
+        consumer = newConsumer(time, new StringDeserializer(), new StringDeserializer());
+        assertThrows(IllegalArgumentException.class, () -> consumer.assign(null));
+    }
+
+    @Test
+    public void testAssignOnEmptyTopicPartition() {
+        consumer = spy(newConsumer(time, new StringDeserializer(), new StringDeserializer()));
+        consumer.assign(Collections.emptyList());
+        assertTrue(consumer.subscription().isEmpty());
+        assertTrue(consumer.assignment().isEmpty());
+    }
+
+    @Test
+    public void testAssignOnNullTopicInPartition() {
+        consumer = newConsumer(time, new StringDeserializer(), new StringDeserializer());
+        assertThrows(IllegalArgumentException.class, () -> consumer.assign(singleton(new TopicPartition(null, 0))));
+    }
+
+    @Test
+    public void testAssignOnEmptyTopicInPartition() {
         consumer = newConsumer(time, new StringDeserializer(), new StringDeserializer());
-        assertThrows(KafkaException.class, consumer::assignment, "not implemented exception");
+        assertThrows(IllegalArgumentException.class, () -> consumer.assign(singleton(new TopicPartition("  ", 0))));
     }
 
     private HashMap<TopicPartition, OffsetAndMetadata> mockTopicPartitionOffset() {