KAFKA-13763 (2): Refactor IncrementalCooperativeAssignor for improved unit testing

connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ConnectAssignor.java

@@ -32,7 +32,7 @@ public interface ConnectAssignor {
      * method computes an assignment of connectors and tasks among the members of the worker group.
      *
      * @param leaderId the leader of the group
-     * @param protocol the protocol type; for Connect assignors this is normally "connect"
+     * @param protocol the protocol type; for Connect assignors this is "eager", "compatible", or "sessioned"
      * @param allMemberMetadata the metadata of all the active workers of the group
      * @param coordinator the worker coordinator that runs this assignor
      * @return the assignment of connectors and tasks to workers

connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/IncrementalCooperativeConnectProtocol.java

@@ -154,7 +154,7 @@ public static ByteBuffer serializeMetadata(ExtendedWorkerState workerState, bool
                 .set(CONFIG_OFFSET_KEY_NAME, workerState.offset());
         // Not a big issue if we embed the protocol version with the assignment in the metadata
         Struct allocation = new Struct(ALLOCATION_V1)
-                .set(ALLOCATION_KEY_NAME, serializeAssignment(workerState.assignment()));
+                .set(ALLOCATION_KEY_NAME, serializeAssignment(workerState.assignment(), sessioned));
         Struct connectProtocolHeader = sessioned ? CONNECT_PROTOCOL_HEADER_V2 : CONNECT_PROTOCOL_HEADER_V1;
         ByteBuffer buffer = ByteBuffer.allocate(connectProtocolHeader.sizeOf()
                                                 + CONFIG_STATE_V1.sizeOf(configState)
@@ -230,15 +230,16 @@ public static ExtendedWorkerState deserializeMetadata(ByteBuffer buffer) {
      *   ScheduledDelay     => Int32
      * </pre>
      */
-    public static ByteBuffer serializeAssignment(ExtendedAssignment assignment) {
+    public static ByteBuffer serializeAssignment(ExtendedAssignment assignment, boolean sessioned) {
         // comparison depends on reference equality for now
         if (assignment == null || ExtendedAssignment.empty().equals(assignment)) {
             return null;
         }
         Struct struct = assignment.toStruct();
-        ByteBuffer buffer = ByteBuffer.allocate(CONNECT_PROTOCOL_HEADER_V1.sizeOf()
+        Struct protocolHeader = sessioned ? CONNECT_PROTOCOL_HEADER_V2 : CONNECT_PROTOCOL_HEADER_V1;
+        ByteBuffer buffer = ByteBuffer.allocate(protocolHeader.sizeOf()
                                                 + ASSIGNMENT_V1.sizeOf(struct));
-        CONNECT_PROTOCOL_HEADER_V1.writeTo(buffer);
+        protocolHeader.writeTo(buffer);
         ASSIGNMENT_V1.write(buffer, struct);
         buffer.flip();
         return buffer;

connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/WorkerCoordinator.java

@@ -416,6 +416,29 @@ private String ownerUrl(String connector) {
             return allMembers.get(ownerId).url();
         }
 
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (!(o instanceof LeaderState)) return false;
+            LeaderState that = (LeaderState) o;
+            return Objects.equals(allMembers, that.allMembers)
+                    && Objects.equals(connectorOwners, that.connectorOwners)
+                    && Objects.equals(taskOwners, that.taskOwners);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(allMembers, connectorOwners, taskOwners);
+        }
+
+        @Override
+        public String toString() {
+            return "LeaderState{"
+                    + "allMembers=" + allMembers
+                    + ", connectorOwners=" + connectorOwners
+                    + ", taskOwners=" + taskOwners
+                    + '}';
+        }
     }
 
     public static class ConnectorsAndTasks {

connect/runtime/src/main/java/org/apache/kafka/connect/util/ConnectUtils.java

@@ -30,10 +30,15 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.Collection;
+import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.concurrent.ExecutionException;
+import java.util.function.Function;
+import java.util.stream.Collector;
+import java.util.stream.Collectors;
 
 public final class ConnectUtils {
     private static final Logger log = LoggerFactory.getLogger(ConnectUtils.class);
@@ -160,4 +165,30 @@ public static boolean isSourceConnector(Connector connector) {
         return SourceConnector.class.isAssignableFrom(connector.getClass());
     }
 
+    public static <K, I, O> Map<K, O> transformValues(Map<K, I> map, Function<I, O> transformation) {
+        return map.entrySet().stream().collect(Collectors.toMap(
+                Map.Entry::getKey,
+                transformation.compose(Map.Entry::getValue)
+        ));
+    }
+
+    public static <I> List<I> combineCollections(Collection<Collection<I>> collections) {
+        return combineCollections(collections, Function.identity());
+    }
+
+    public static <I, T> List<T> combineCollections(Collection<I> collection, Function<I, Collection<T>> extractCollection) {
+        return combineCollections(collection, extractCollection, Collectors.toList());
+    }
+
+    public static <I, T, C> C combineCollections(
+            Collection<I> collection,
+            Function<I, Collection<T>> extractCollection,
+            Collector<T, ?, C> collector
+    ) {
+        return collection.stream()
+                .map(extractCollection)
+                .flatMap(Collection::stream)
+                .collect(collector);
+    }
+
 }

connect/runtime/src/test/java/org/apache/kafka/connect/runtime/distributed/ConnectProtocolCompatibilityTest.java

@@ -16,32 +16,22 @@
  */
 package org.apache.kafka.connect.runtime.distributed;
 
-import org.apache.kafka.connect.runtime.TargetState;
-import org.apache.kafka.connect.storage.KafkaConfigBackingStore;
 import org.apache.kafka.connect.util.ConnectorTaskId;
-import org.junit.After;
-import org.junit.Before;
-import org.junit.Rule;
 import org.junit.Test;
-import org.mockito.Mock;
-import org.mockito.junit.MockitoJUnit;
-import org.mockito.junit.MockitoRule;
 
 import java.nio.ByteBuffer;
 import java.util.Arrays;
 import java.util.Collections;
-import java.util.HashMap;
 
 import static org.apache.kafka.connect.runtime.distributed.IncrementalCooperativeConnectProtocol.CONNECT_PROTOCOL_V1;
+import static org.apache.kafka.connect.runtime.distributed.IncrementalCooperativeConnectProtocol.CONNECT_PROTOCOL_V2;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.verifyNoMoreInteractions;
-import static org.mockito.Mockito.when;
 
 public class ConnectProtocolCompatibilityTest {
+    private static final String LEADER = "leader";
     private static final String LEADER_URL = "leaderUrl:8083";
+    private static final long CONFIG_OFFSET = 1;
 
     private String connectorId1 = "connector1";
     private String connectorId2 = "connector2";
@@ -51,95 +41,58 @@ public class ConnectProtocolCompatibilityTest {
     private ConnectorTaskId taskId2x0 = new ConnectorTaskId(connectorId2, 0);
     private ConnectorTaskId taskId3x0 = new ConnectorTaskId(connectorId3, 0);
 
-    @Rule
-    public MockitoRule rule = MockitoJUnit.rule();
-
-    @Mock
-    private KafkaConfigBackingStore configStorage;
-    private ClusterConfigState configState;
-
-    @Before
-    public void setup() {
-        configStorage = mock(KafkaConfigBackingStore.class);
-        configState = new ClusterConfigState(
-                1L,
-                null,
-                Collections.singletonMap(connectorId1, 1),
-                Collections.singletonMap(connectorId1, new HashMap<>()),
-                Collections.singletonMap(connectorId1, TargetState.STARTED),
-                Collections.singletonMap(taskId1x0, new HashMap<>()),
-                Collections.emptySet());
-    }
-
-    @After
-    public void teardown() {
-        verifyNoMoreInteractions(configStorage);
-    }
-
     @Test
     public void testEagerToEagerMetadata() {
-        when(configStorage.snapshot()).thenReturn(configState);
-        ExtendedWorkerState workerState = new ExtendedWorkerState(LEADER_URL, configStorage.snapshot().offset(), null);
+        ConnectProtocol.WorkerState workerState = emptyWorkerState();
         ByteBuffer metadata = ConnectProtocol.serializeMetadata(workerState);
         ConnectProtocol.WorkerState state = ConnectProtocol.deserializeMetadata(metadata);
         assertEquals(LEADER_URL, state.url());
         assertEquals(1, state.offset());
-        verify(configStorage).snapshot();
     }
 
     @Test
     public void testCoopToCoopMetadata() {
-        when(configStorage.snapshot()).thenReturn(configState);
-        ExtendedWorkerState workerState = new ExtendedWorkerState(LEADER_URL, configStorage.snapshot().offset(), null);
+        ExtendedWorkerState workerState = emptyExtendedWorkerState(CONNECT_PROTOCOL_V1);
         ByteBuffer metadata = IncrementalCooperativeConnectProtocol.serializeMetadata(workerState, false);
         ExtendedWorkerState state = IncrementalCooperativeConnectProtocol.deserializeMetadata(metadata);
         assertEquals(LEADER_URL, state.url());
         assertEquals(1, state.offset());
-        verify(configStorage).snapshot();
     }
 
     @Test
     public void testSessionedToCoopMetadata() {
-        when(configStorage.snapshot()).thenReturn(configState);
-        ExtendedWorkerState workerState = new ExtendedWorkerState(LEADER_URL, configStorage.snapshot().offset(), null);
+        ExtendedWorkerState workerState = emptyExtendedWorkerState(CONNECT_PROTOCOL_V2);
         ByteBuffer metadata = IncrementalCooperativeConnectProtocol.serializeMetadata(workerState, true);
         ExtendedWorkerState state = IncrementalCooperativeConnectProtocol.deserializeMetadata(metadata);
         assertEquals(LEADER_URL, state.url());
         assertEquals(1, state.offset());
-        verify(configStorage).snapshot();
     }
 
     @Test
     public void testSessionedToEagerMetadata() {
-        when(configStorage.snapshot()).thenReturn(configState);
-        ExtendedWorkerState workerState = new ExtendedWorkerState(LEADER_URL, configStorage.snapshot().offset(), null);
+        ExtendedWorkerState workerState = emptyExtendedWorkerState(CONNECT_PROTOCOL_V2);
         ByteBuffer metadata = IncrementalCooperativeConnectProtocol.serializeMetadata(workerState, true);
         ConnectProtocol.WorkerState state = ConnectProtocol.deserializeMetadata(metadata);
         assertEquals(LEADER_URL, state.url());
         assertEquals(1, state.offset());
-        verify(configStorage).snapshot();
     }
 
     @Test
     public void testCoopToEagerMetadata() {
-        when(configStorage.snapshot()).thenReturn(configState);
-        ExtendedWorkerState workerState = new ExtendedWorkerState(LEADER_URL, configStorage.snapshot().offset(), null);
+        ExtendedWorkerState workerState = emptyExtendedWorkerState(CONNECT_PROTOCOL_V1);
         ByteBuffer metadata = IncrementalCooperativeConnectProtocol.serializeMetadata(workerState, false);
         ConnectProtocol.WorkerState state = ConnectProtocol.deserializeMetadata(metadata);
         assertEquals(LEADER_URL, state.url());
         assertEquals(1, state.offset());
-        verify(configStorage).snapshot();
     }
 
     @Test
     public void testEagerToCoopMetadata() {
-        when(configStorage.snapshot()).thenReturn(configState);
-        ConnectProtocol.WorkerState workerState = new ConnectProtocol.WorkerState(LEADER_URL, configStorage.snapshot().offset());
+        ConnectProtocol.WorkerState workerState = emptyWorkerState();
         ByteBuffer metadata = ConnectProtocol.serializeMetadata(workerState);
         ConnectProtocol.WorkerState state = IncrementalCooperativeConnectProtocol.deserializeMetadata(metadata);
         assertEquals(LEADER_URL, state.url());
         assertEquals(1, state.offset());
-        verify(configStorage).snapshot();
     }
 
     @Test
@@ -176,7 +129,7 @@ public void testCoopToCoopAssignment() {
                 Arrays.asList(connectorId1, connectorId3), Arrays.asList(taskId2x0),
                 Collections.emptyList(), Collections.emptyList(), 0);
 
-        ByteBuffer leaderBuf = IncrementalCooperativeConnectProtocol.serializeAssignment(assignment);
+        ByteBuffer leaderBuf = IncrementalCooperativeConnectProtocol.serializeAssignment(assignment, false);
         ConnectProtocol.Assignment leaderAssignment = ConnectProtocol.deserializeAssignment(leaderBuf);
         assertFalse(leaderAssignment.failed());
         assertEquals("leader", leaderAssignment.leader());
@@ -235,7 +188,7 @@ public void testCoopToEagerAssignment() {
                 Arrays.asList(connectorId1, connectorId3), Arrays.asList(taskId2x0),
                 Collections.emptyList(), Collections.emptyList(), 0);
 
-        ByteBuffer leaderBuf = IncrementalCooperativeConnectProtocol.serializeAssignment(assignment);
+        ByteBuffer leaderBuf = IncrementalCooperativeConnectProtocol.serializeAssignment(assignment, false);
         ConnectProtocol.Assignment leaderAssignment = ConnectProtocol.deserializeAssignment(leaderBuf);
         assertFalse(leaderAssignment.failed());
         assertEquals("leader", leaderAssignment.leader());
@@ -248,7 +201,7 @@ public void testCoopToEagerAssignment() {
                 Arrays.asList(connectorId2), Arrays.asList(taskId1x0, taskId3x0),
                 Collections.emptyList(), Collections.emptyList(), 0);
 
-        ByteBuffer memberBuf = IncrementalCooperativeConnectProtocol.serializeAssignment(assignment2);
+        ByteBuffer memberBuf = IncrementalCooperativeConnectProtocol.serializeAssignment(assignment2, false);
         ConnectProtocol.Assignment memberAssignment = ConnectProtocol.deserializeAssignment(memberBuf);
         assertFalse(memberAssignment.failed());
         assertEquals("member", memberAssignment.leader());
@@ -257,4 +210,24 @@ public void testCoopToEagerAssignment() {
         assertEquals(Arrays.asList(taskId1x0, taskId3x0), memberAssignment.tasks());
     }
 
+    private ConnectProtocol.WorkerState emptyWorkerState() {
+        return new ConnectProtocol.WorkerState(LEADER_URL, CONFIG_OFFSET);
+    }
+
+    private ExtendedWorkerState emptyExtendedWorkerState(short protocolVersion) {
+        ExtendedAssignment assignment = new ExtendedAssignment(
+                protocolVersion,
+                ConnectProtocol.Assignment.NO_ERROR,
+                LEADER,
+                LEADER_URL,
+                CONFIG_OFFSET,
+                Collections.emptySet(),
+                Collections.emptySet(),
+                Collections.emptySet(),
+                Collections.emptySet(),
+                0
+        );
+        return new ExtendedWorkerState(LEADER_URL, CONFIG_OFFSET, assignment);
+    }
+
 }