StreamThread extends Thread { public Map<TaskId, StreamTask> tasks() { return Collections.unmodifiableMap(activeTasks); } StreamThread(TopologyBuilder builder,
                        StreamsConfig config,
                        KafkaClientSupplier clientSupplier,
                        String applicationId,
                        String clientId,
                        UUID processId,
                        Metrics metrics,
                        Time time); void partitionAssignor(StreamPartitionAssignor partitionAssignor); @Override void run(); void close(); Map<TaskId, StreamTask> tasks(); Set<TaskId> prevTasks(); Set<TaskId> cachedTasks(); final PartitionGrouper partitionGrouper; final String applicationId; final String clientId; final UUID processId; }
@SuppressWarnings("unchecked") @Test public void testPartitionAssignmentChange() throws Exception { StreamsConfig config = new StreamsConfig(configProps()); TopologyBuilder builder = new TopologyBuilder(); builder.addSource("source1", "topic1"); builder.addSource("source2", "topic2"); builder.addSource("source3", "topic3"); builder.addProcessor("processor", new MockProcessorSupplier(), "source2", "source3"); StreamThread thread = new StreamThread(builder, config, new MockClientSupplier(), applicationId, clientId, processId, new Metrics(), new SystemTime()) { @Override protected StreamTask createStreamTask(TaskId id, Collection<TopicPartition> partitionsForTask) { ProcessorTopology topology = builder.build("X", id.topicGroupId); return new TestStreamTask(id, applicationId, partitionsForTask, topology, consumer, producer, restoreConsumer, config); } }; initPartitionGrouper(config, thread); ConsumerRebalanceListener rebalanceListener = thread.rebalanceListener; assertTrue(thread.tasks().isEmpty()); List<TopicPartition> revokedPartitions; List<TopicPartition> assignedPartitions; Set<TopicPartition> expectedGroup1; Set<TopicPartition> expectedGroup2; revokedPartitions = Collections.emptyList(); assignedPartitions = Collections.singletonList(t1p1); expectedGroup1 = new HashSet<>(Arrays.asList(t1p1)); rebalanceListener.onPartitionsRevoked(revokedPartitions); rebalanceListener.onPartitionsAssigned(assignedPartitions); assertTrue(thread.tasks().containsKey(task1)); assertEquals(expectedGroup1, thread.tasks().get(task1).partitions()); assertEquals(1, thread.tasks().size()); revokedPartitions = assignedPartitions; assignedPartitions = Collections.singletonList(t1p2); expectedGroup2 = new HashSet<>(Arrays.asList(t1p2)); rebalanceListener.onPartitionsRevoked(revokedPartitions); rebalanceListener.onPartitionsAssigned(assignedPartitions); assertTrue(thread.tasks().containsKey(task2)); assertEquals(expectedGroup2, thread.tasks().get(task2).partitions()); assertEquals(1, thread.tasks().size()); revokedPartitions = assignedPartitions; assignedPartitions = Arrays.asList(t1p1, t1p2); expectedGroup1 = new HashSet<>(Collections.singleton(t1p1)); expectedGroup2 = new HashSet<>(Collections.singleton(t1p2)); rebalanceListener.onPartitionsRevoked(revokedPartitions); rebalanceListener.onPartitionsAssigned(assignedPartitions); assertTrue(thread.tasks().containsKey(task1)); assertTrue(thread.tasks().containsKey(task2)); assertEquals(expectedGroup1, thread.tasks().get(task1).partitions()); assertEquals(expectedGroup2, thread.tasks().get(task2).partitions()); assertEquals(2, thread.tasks().size()); revokedPartitions = assignedPartitions; assignedPartitions = Arrays.asList(t2p1, t2p2, t3p1, t3p2); expectedGroup1 = new HashSet<>(Arrays.asList(t2p1, t3p1)); expectedGroup2 = new HashSet<>(Arrays.asList(t2p2, t3p2)); rebalanceListener.onPartitionsRevoked(revokedPartitions); rebalanceListener.onPartitionsAssigned(assignedPartitions); assertTrue(thread.tasks().containsKey(task4)); assertTrue(thread.tasks().containsKey(task5)); assertEquals(expectedGroup1, thread.tasks().get(task4).partitions()); assertEquals(expectedGroup2, thread.tasks().get(task5).partitions()); assertEquals(2, thread.tasks().size()); revokedPartitions = assignedPartitions; assignedPartitions = Arrays.asList(t1p1, t2p1, t3p1); expectedGroup1 = new HashSet<>(Arrays.asList(t1p1)); expectedGroup2 = new HashSet<>(Arrays.asList(t2p1, t3p1)); rebalanceListener.onPartitionsRevoked(revokedPartitions); rebalanceListener.onPartitionsAssigned(assignedPartitions); assertTrue(thread.tasks().containsKey(task1)); assertTrue(thread.tasks().containsKey(task4)); assertEquals(expectedGroup1, thread.tasks().get(task1).partitions()); assertEquals(expectedGroup2, thread.tasks().get(task4).partitions()); assertEquals(2, thread.tasks().size()); revokedPartitions = assignedPartitions; assignedPartitions = Collections.emptyList(); rebalanceListener.onPartitionsRevoked(revokedPartitions); rebalanceListener.onPartitionsAssigned(assignedPartitions); assertTrue(thread.tasks().isEmpty()); }