StreamPartitionAssignor implements PartitionAssignor, Configurable { @Override public void onAssignment(Assignment assignment) { List<TopicPartition> partitions = new ArrayList<>(assignment.partitions()); Collections.sort(partitions, PARTITION_COMPARATOR); AssignmentInfo info = AssignmentInfo.decode(assignment.userData()); this.standbyTasks = info.standbyTasks; Map<TopicPartition, Set<TaskId>> partitionToTaskIds = new HashMap<>(); Iterator<TaskId> iter = info.activeTasks.iterator(); for (TopicPartition partition : partitions) { Set<TaskId> taskIds = partitionToTaskIds.get(partition); if (taskIds == null) { taskIds = new HashSet<>(); partitionToTaskIds.put(partition, taskIds); } if (iter.hasNext()) { taskIds.add(iter.next()); } else { TaskAssignmentException ex = new TaskAssignmentException( "failed to find a task id for the partition=" + partition.toString() + ", partitions=" + partitions.size() + ", assignmentInfo=" + info.toString() ); log.error(ex.getMessage(), ex); throw ex; } } this.partitionToTaskIds = partitionToTaskIds; }  @Override void configure(Map<String, ?> configs); @Override String name(); @Override Subscription subscription(Set<String> topics); @Override Map<String, Assignment> assign(Cluster metadata, Map<String, Subscription> subscriptions); @Override void onAssignment(Assignment assignment); Set<TaskId> tasksForState(String stateName); Set<TaskId> tasksForPartition(TopicPartition partition); Map<TaskId, Set<TopicPartition>> standbyTasks(); void setInternalTopicManager(InternalTopicManager internalTopicManager);  }
@Test public void testOnAssignment() throws Exception { StreamsConfig config = new StreamsConfig(configProps()); TopicPartition t2p3 = new TopicPartition("topic2", 3); TopologyBuilder builder = new TopologyBuilder(); builder.addSource("source1", "topic1"); builder.addSource("source2", "topic2"); builder.addProcessor("processor", new MockProcessorSupplier(), "source1", "source2"); UUID uuid = UUID.randomUUID(); String client1 = "client1"; StreamThread thread = new StreamThread(builder, config, new MockClientSupplier(), "test", client1, uuid, new Metrics(), new SystemTime()); StreamPartitionAssignor partitionAssignor = new StreamPartitionAssignor(); partitionAssignor.configure(config.getConsumerConfigs(thread, "test", client1)); List<TaskId> activeTaskList = Utils.mkList(task0, task3); Map<TaskId, Set<TopicPartition>> standbyTasks = new HashMap<>(); standbyTasks.put(task1, Utils.mkSet(new TopicPartition("t1", 0))); standbyTasks.put(task2, Utils.mkSet(new TopicPartition("t2", 0))); AssignmentInfo info = new AssignmentInfo(activeTaskList, standbyTasks); PartitionAssignor.Assignment assignment = new PartitionAssignor.Assignment(Utils.mkList(t1p0, t2p3), info.encode()); partitionAssignor.onAssignment(assignment); assertEquals(Utils.mkSet(task0), partitionAssignor.tasksForPartition(t1p0)); assertEquals(Utils.mkSet(task3), partitionAssignor.tasksForPartition(t2p3)); assertEquals(standbyTasks, partitionAssignor.standbyTasks()); }