org.apache.kafka.clients.admin.CreatePartitionsResult Java Examples

@Test
public void testIncreasePartitions() {
  TopicService service = new TopicServiceImpl(adminClient, true);
  CreatePartitionsResult result = mock(CreatePartitionsResult.class);
  when(result.all()).thenReturn(KafkaFuture.completedFuture(null));
  when(adminClient.createPartitions(any(Map.class), any(CreatePartitionsOptions.class))).thenReturn(result);

  service.increasePartitions(Collections.singletonList(
      new ConfiguredTopic("test", 3, (short) 2, Collections.emptyMap())));

  ArgumentCaptor<Map> increase = ArgumentCaptor.forClass(Map.class);
  ArgumentCaptor<CreatePartitionsOptions> options = ArgumentCaptor.forClass(CreatePartitionsOptions.class);
  verify(adminClient).createPartitions((Map<String, NewPartitions>) increase.capture(), options.capture());
  Assert.assertEquals(1, increase.getValue().size());
  Assert.assertTrue(increase.getValue().containsKey("test"));
  Assert.assertEquals(3, ((NewPartitions) increase.getValue().get("test")).totalCount());
  Assert.assertTrue(options.getValue().validateOnly());
}
 

@Test
public void testMaybeIncreasePartitionCount() throws InterruptedException, ExecutionException, TimeoutException {
  AdminClient adminClient = EasyMock.createMock(AdminClient.class);
  NewTopic topicToAddPartitions = SamplingUtils.wrapTopic(MOCK_TOPIC, MOCK_DESIRED_PARTITION_COUNT,
                                                          MOCK_REPLICATION_FACTOR, MOCK_DESIRED_RETENTION_MS);
  DescribeTopicsResult describeTopicsResult = EasyMock.createMock(DescribeTopicsResult.class);
  KafkaFuture<TopicDescription> topicDescriptionFuture = EasyMock.createMock(KafkaFuture.class);
  TopicDescription topicDescription = EasyMock.createMock(TopicDescription.class);
  Map<String, KafkaFuture<TopicDescription>> describeTopicsValues = Collections.singletonMap(MOCK_TOPIC, topicDescriptionFuture);
  Map<String, KafkaFuture<Void>> createPartitionsValues = Collections.singletonMap(MOCK_TOPIC, EasyMock.createMock(KafkaFuture.class));
  CreatePartitionsResult createPartitionsResult = EasyMock.createMock(CreatePartitionsResult.class);

  EasyMock.expect(adminClient.describeTopics(Collections.singletonList(MOCK_TOPIC))).andReturn(describeTopicsResult);
  EasyMock.expect(describeTopicsResult.values()).andReturn(describeTopicsValues);
  EasyMock.expect(topicDescriptionFuture.get(CLIENT_REQUEST_TIMEOUT_MS, TimeUnit.MILLISECONDS)).andReturn(topicDescription);
  EasyMock.expect(topicDescription.partitions()).andReturn(MOCK_PARTITIONS);
  EasyMock.expect(adminClient.createPartitions(Collections.singletonMap(MOCK_TOPIC, EasyMock.anyObject())))
          .andReturn(createPartitionsResult);
  EasyMock.expect(createPartitionsResult.values()).andReturn(createPartitionsValues);

  EasyMock.replay(adminClient, describeTopicsResult, topicDescriptionFuture, topicDescription, createPartitionsResult);
  boolean increasePartitionCount = SamplingUtils.maybeIncreasePartitionCount(adminClient, topicToAddPartitions);

  EasyMock.verify(adminClient, describeTopicsResult, topicDescriptionFuture, topicDescription, createPartitionsResult);
  assertTrue(increasePartitionCount);
}
 

protected void alterTopicNumPartitions(String topicName, String resourceName) throws InterruptedException, ExecutionException, TimeoutException {
    int changedValue = 2;

    NewPartitions newPartitions = NewPartitions.increaseTo(changedValue);
    Map<String, NewPartitions> map = new HashMap<>(1);
    map.put(topicName, newPartitions);

    CreatePartitionsResult createPartitionsResult = adminClient.createPartitions(map);
    createPartitionsResult.all().get();

    // Wait for the resource to be modified
    waitFor(() -> {
        KafkaTopic topic = operation().inNamespace(NAMESPACE).withName(resourceName).get();
        LOGGER.info("Polled topic {}, waiting for partitions change", resourceName);
        int gotValue = TopicSerialization.fromTopicResource(topic).getNumPartitions();
        LOGGER.info("Expected value {}, got value {}", changedValue, gotValue);
        return changedValue == gotValue;
    }, "Expected the topic " + topicName + "to have " + changedValue + " partitions by now");
}
 

/**
 * 给topic增加分区
 */
public CreatePartitionsResult addPartitionsForTopic(String topic, int partitions) {
	Map<String, NewPartitions> map = new HashMap<>();
	NewPartitions np = NewPartitions.increaseTo(partitions);
	map.put(topic, np);
	CreatePartitionsOptions cpo = new CreatePartitionsOptions();
	cpo.timeoutMs(5 * 1000);
	return adminClient.createPartitions(map, cpo);
}
 

void maybeAddPartitions(int minPartitionNum) throws ExecutionException, InterruptedException {
  Map<String, KafkaFuture<TopicDescription>> kafkaFutureMap =
      _adminClient.describeTopics(Collections.singleton(_topic)).values();
  KafkaFuture<TopicDescription> topicDescriptions = kafkaFutureMap.get(_topic);
  List<TopicPartitionInfo> partitions = topicDescriptions.get().partitions();

  int partitionNum = partitions.size();
  if (partitionNum < minPartitionNum) {
    LOGGER.info("{} will increase partition of the topic {} in the cluster from {}"
        + " to {}.", this.getClass().toString(), _topic, partitionNum, minPartitionNum);
    Set<Integer> blackListedBrokers = _topicFactory.getBlackListedBrokers(_zkConnect);
    Set<BrokerMetadata> brokers = new HashSet<>();
    for (Node broker : _adminClient.describeCluster().nodes().get()) {
      BrokerMetadata brokerMetadata = new BrokerMetadata(
          broker.id(), null
      );
      brokers.add(brokerMetadata);
    }

    if (!blackListedBrokers.isEmpty()) {
      brokers.removeIf(broker -> blackListedBrokers.contains(broker.id()));
    }

    List<List<Integer>> newPartitionAssignments = newPartitionAssignments(minPartitionNum, partitionNum, brokers, _replicationFactor);

    NewPartitions newPartitions = NewPartitions.increaseTo(minPartitionNum, newPartitionAssignments);

    Map<String, NewPartitions> newPartitionsMap = new HashMap<>();
    newPartitionsMap.put(_topic, newPartitions);
    CreatePartitionsResult createPartitionsResult = _adminClient.createPartitions(newPartitionsMap);

  }
}
 

protected CreatePartitionsResult increasePartitionsTo(String topicName, int numPartitions) {
  return adminClient.createPartitions(ImmutableMap.of(topicName, NewPartitions.increaseTo(numPartitions)));
}