/** Simulate a rebalance event. */ public synchronized void rebalance(Collection<TopicPartition> newAssignment) { // TODO: Rebalance callbacks this.records.clear(); this.subscriptions.assignFromSubscribed(newAssignment); }
@Override public synchronized void subscribe(Pattern pattern, final ConsumerRebalanceListener listener) { ensureNotClosed(); committed.clear(); this.subscriptions.subscribe(pattern, listener); Set<String> topicsToSubscribe = new HashSet<>(); for (String topic: partitions.keySet()) { if (pattern.matcher(topic).matches() && !subscriptions.subscription().contains(topic)) topicsToSubscribe.add(topic); } ensureNotClosed(); this.subscriptions.subscribeFromPattern(topicsToSubscribe); final Set<TopicPartition> assignedPartitions = new HashSet<>(); for (final String topic : topicsToSubscribe) { for (final PartitionInfo info : this.partitions.get(topic)) { assignedPartitions.add(new TopicPartition(topic, info.partition())); } } subscriptions.assignFromSubscribed(assignedPartitions); }
@Test(expected = IllegalArgumentException.class) public void cantAssignPartitionForUnmatchedPattern() { state.subscribe(Pattern.compile(".*t"), rebalanceListener); state.subscribeFromPattern(new HashSet<>(Collections.singletonList(topic))); state.assignFromSubscribed(Collections.singletonList(t1p0)); }
@Test(expected = IllegalArgumentException.class) public void cantAssignPartitionForUnsubscribedTopics() { state.subscribe(singleton(topic), rebalanceListener); state.assignFromSubscribed(Collections.singletonList(t1p0)); }
@Test(expected = IllegalStateException.class) public void invalidPositionUpdate() { state.subscribe(singleton(topic), rebalanceListener); state.assignFromSubscribed(singleton(tp0)); state.position(tp0, 0); }
subscriptions.assignFromSubscribed(assignment.partitions());
@Test public void testFetchDuringRebalance() { subscriptions.subscribe(singleton(topicName), listener); subscriptions.assignFromSubscribed(singleton(tp0)); subscriptions.seek(tp0, 0); assertEquals(1, fetcher.sendFetches()); // Now the rebalance happens and fetch positions are cleared subscriptions.assignFromSubscribed(singleton(tp0)); client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.NONE, 100L, 0)); consumerClient.poll(time.timer(0)); // The active fetch should be ignored since its position is no longer valid assertTrue(fetcher.fetchedRecords().isEmpty()); }
@Test public void verifyAssignmentListener() { final AtomicReference<Set<TopicPartition>> assignmentRef = new AtomicReference<>(); state.addListener(new SubscriptionState.Listener() { @Override public void onAssignment(Set<TopicPartition> assignment) { assignmentRef.set(assignment); } }); Set<TopicPartition> userAssignment = Utils.mkSet(tp0, tp1); state.assignFromUser(userAssignment); assertEquals(userAssignment, assignmentRef.get()); state.unsubscribe(); assertEquals(Collections.emptySet(), assignmentRef.get()); Set<TopicPartition> autoAssignment = Utils.mkSet(t1p0); state.subscribe(singleton(topic1), rebalanceListener); state.assignFromSubscribed(autoAssignment); assertEquals(autoAssignment, assignmentRef.get()); }
@Test public void topicSubscription() { state.subscribe(singleton(topic), rebalanceListener); assertEquals(1, state.subscription().size()); assertTrue(state.assignedPartitions().isEmpty()); assertEquals(0, state.numAssignedPartitions()); assertTrue(state.partitionsAutoAssigned()); state.assignFromSubscribed(singleton(tp0)); state.seek(tp0, 1); assertEquals(1L, state.position(tp0).longValue()); state.assignFromSubscribed(singleton(tp1)); assertTrue(state.isAssigned(tp1)); assertFalse(state.isAssigned(tp0)); assertFalse(state.isFetchable(tp1)); assertEquals(singleton(tp1), state.assignedPartitions()); assertEquals(1, state.numAssignedPartitions()); }
@Test public void unsubscription() { state.subscribe(Pattern.compile(".*"), rebalanceListener); state.subscribeFromPattern(new HashSet<>(Arrays.asList(topic, topic1))); state.assignFromSubscribed(singleton(tp1)); assertEquals(singleton(tp1), state.assignedPartitions()); assertEquals(1, state.numAssignedPartitions()); state.unsubscribe(); assertEquals(0, state.subscription().size()); assertTrue(state.assignedPartitions().isEmpty()); assertEquals(0, state.numAssignedPartitions()); state.assignFromUser(singleton(tp0)); assertEquals(singleton(tp0), state.assignedPartitions()); assertEquals(1, state.numAssignedPartitions()); state.unsubscribe(); assertEquals(0, state.subscription().size()); assertTrue(state.assignedPartitions().isEmpty()); assertEquals(0, state.numAssignedPartitions()); }
assertEquals(0, state.numAssignedPartitions()); state.assignFromSubscribed(singleton(tp1)); assertEquals(singleton(topic), state.subscription()); state.assignFromSubscribed(Collections.singletonList(t1p0)); assertEquals(1, state.numAssignedPartitions()); state.assignFromSubscribed(Collections.singletonList(tp0));
@Test public void partitionAssignmentChangeOnTopicSubscription() { state.assignFromUser(new HashSet<>(Arrays.asList(tp0, tp1))); // assigned partitions should immediately change assertEquals(2, state.assignedPartitions().size()); assertEquals(2, state.numAssignedPartitions()); assertTrue(state.assignedPartitions().contains(tp0)); assertTrue(state.assignedPartitions().contains(tp1)); state.unsubscribe(); // assigned partitions should immediately change assertTrue(state.assignedPartitions().isEmpty()); assertEquals(0, state.numAssignedPartitions()); state.subscribe(singleton(topic1), rebalanceListener); // assigned partitions should remain unchanged assertTrue(state.assignedPartitions().isEmpty()); assertEquals(0, state.numAssignedPartitions()); state.assignFromSubscribed(singleton(t1p0)); // assigned partitions should immediately change assertEquals(singleton(t1p0), state.assignedPartitions()); assertEquals(1, state.numAssignedPartitions()); state.subscribe(singleton(topic), rebalanceListener); // assigned partitions should remain unchanged assertEquals(singleton(t1p0), state.assignedPartitions()); assertEquals(1, state.numAssignedPartitions()); state.unsubscribe(); // assigned partitions should immediately change assertTrue(state.assignedPartitions().isEmpty()); assertEquals(0, state.numAssignedPartitions()); }
@Test public void testUnknownConsumerId() { client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE)); coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE)); // illegal_generation will cause re-partition subscriptions.subscribe(singleton(topic1), rebalanceListener); subscriptions.assignFromSubscribed(Collections.singletonList(t1p)); time.sleep(sessionTimeoutMs); RequestFuture<Void> future = coordinator.sendHeartbeatRequest(); // should send out the heartbeat assertEquals(1, consumerClient.pendingRequestCount()); assertFalse(future.isDone()); client.prepareResponse(heartbeatResponse(Errors.UNKNOWN_MEMBER_ID)); time.sleep(sessionTimeoutMs); consumerClient.poll(time.timer(0)); assertTrue(future.isDone()); assertTrue(future.failed()); assertEquals(Errors.UNKNOWN_MEMBER_ID.exception(), future.exception()); assertTrue(coordinator.rejoinNeededOrPending()); }
@Test public void testIllegalGeneration() { client.prepareResponse(groupCoordinatorResponse(node, Errors.NONE)); coordinator.ensureCoordinatorReady(time.timer(Long.MAX_VALUE)); // illegal_generation will cause re-partition subscriptions.subscribe(singleton(topic1), rebalanceListener); subscriptions.assignFromSubscribed(Collections.singletonList(t1p)); time.sleep(sessionTimeoutMs); RequestFuture<Void> future = coordinator.sendHeartbeatRequest(); // should send out the heartbeat assertEquals(1, consumerClient.pendingRequestCount()); assertFalse(future.isDone()); client.prepareResponse(heartbeatResponse(Errors.ILLEGAL_GENERATION)); time.sleep(sessionTimeoutMs); consumerClient.poll(time.timer(0)); assertTrue(future.isDone()); assertTrue(future.failed()); assertEquals(Errors.ILLEGAL_GENERATION.exception(), future.exception()); assertTrue(coordinator.rejoinNeededOrPending()); }