Java Code Examples for org.apache.kafka.clients.consumer.Consumer#assign()

public NakadiKafkaConsumer(
        final Consumer<byte[], byte[]> kafkaConsumer,
        final List<KafkaCursor> kafkaCursors,
        final Map<TopicPartition, Timeline> timelineMap,
        final long pollTimeout) {
    this.kafkaConsumer = kafkaConsumer;
    this.pollTimeout = pollTimeout;
    this.timelineMap = timelineMap;
    // define topic/partitions to consume from
    final Map<TopicPartition, KafkaCursor> topicCursors = kafkaCursors.stream().collect(
            Collectors.toMap(
                    cursor -> new TopicPartition(cursor.getTopic(), cursor.getPartition()),
                    cursor -> cursor,
                    (cursor1, cursor2) -> cursor2
            ));
    kafkaConsumer.assign(new ArrayList<>(topicCursors.keySet()));
    topicCursors.forEach((topicPartition, cursor) -> kafkaConsumer.seek(topicPartition, cursor.getOffset()));
}
 

/**
 * createConsumer - create a new consumer
 *
 * @return
 * @throws Exception
 */
private Consumer<String, String> createConsumer() throws Exception {

    // Seek to end automatically
    TopicPartition dataTopicPartition = new TopicPartition(topicName, 0);
    List<TopicPartition> topics = Arrays.asList(dataTopicPartition);

    Properties props = ConfUtils.getProps(CONSUMER_PROPS);
    Consumer<String, String> consumer = new KafkaConsumer<>(props);
    consumer.assign(topics);

    if (offset == -1) {
        consumer.seekToEnd(topics);
        logger.info("Consumer seek to end");
    } else {
        consumer.seek(dataTopicPartition, offset);
        logger.info(String.format("read changed as offset: %s", consumer.position(dataTopicPartition)));
    }
    return consumer;
}
 

static void consumeFail(Consumer<String, String> consumer, String topic) {
    TopicPartition partition = new TopicPartition(topic, 0);
    consumer.assign(Arrays.asList(partition));

    try {
        while (consumer.partitionsFor(topic, Duration.ofSeconds(1)).size() == 0) {
            System.out.println("No assignment yet for consumer");
        }

        consumer.seekToBeginning(Arrays.asList(partition));
        consumer.poll(Duration.ofSeconds(1));

        Assert.fail("Should fail with TopicAuthorizationException");
    } catch (TopicAuthorizationException e) {
    }
}
 

/**
 * Creates a KafkaConsumer that is manually assigned to all partitions of the test topic indicated
 * by the {@code subscription}.
 */
protected Consumer<String, ByteBuffer> getValidationConsumer(String topic, String subscription) {
  Consumer<String, ByteBuffer> consumer =
      kafkaClientFactory.createConsumer(
          ProjectSubscriptionName.of(PROJECT, subscription).toString());

  Set<TopicPartition> topicPartitions =
      consumer
          .listTopics()
          .entrySet()
          .stream()
          .filter(e -> e.getKey().equals(ProjectTopicName.of(PROJECT, topic).toString()))
          .flatMap(
              e -> e.getValue().stream().map(p -> new TopicPartition(p.topic(), p.partition())))
          .collect(Collectors.toSet());
  consumer.assign(topicPartitions);

  return consumer;
}
 

public void syncGroupForOffset(TopicPartition topicPartition, long sourceOffset, long targetOffset) {
    OffsetStoreFile offsetStoreFile = topicOffsets.get(topicPartition);
    // __consumer_offsets contains the offset of the message to read next. So we need to search for the offset + 1
    // if we do not do that we might miss
    List<String> groups = offsetStoreFile.groupForOffset(sourceOffset + 1);
    if (groups != null && groups.size() > 0) {
        for (String group : groups) {
            Map<String, Object> groupConsumerConfig = new HashMap<>(consumerConfig);
            groupConsumerConfig.put("group.id", group);
            Consumer<byte[], byte[]> consumer = new KafkaConsumer<>(groupConsumerConfig);
            consumer.assign(Collections.singletonList(topicPartition));
            // ! Target Offset + 1 as we commit the offset of the "next message to read"
            OffsetAndMetadata offsetAndMetadata = new OffsetAndMetadata(targetOffset + 1);
            Map<TopicPartition, OffsetAndMetadata> offsets = Collections.singletonMap(topicPartition, offsetAndMetadata);
            consumer.commitSync(offsets);
            consumer.close();
            log.debug("Committed target offset {} for group {} for topic {} partition {}",
                    (targetOffset + 1), group, topicPartition.topic(), topicPartition.partition());
        }
    }
}
 

/**
 * Initializes and returns a List of Consumers that are manually assigned to specific
 * TopicPartitions. We choose to use manual assignment to avoid the timeout, blocking, and
 * heartbeats required when using dynamic subscriptions.
 *
 * @return List of Consumers assigned to partitions from the topic
 */
@Override
protected void startUp() {
  // Create first consumer and determine # partitions in topic
  Consumer<String, ByteBuffer> first = kafkaClientFactory.createConsumer(subscription.getName());
  List<PartitionInfo> partitionInfo = first.partitionsFor(kafkaTopicName);

  int totalConsumers = Math.min(partitionInfo.size(), consumerExecutors);
  logger.atFine().log(
      "Assigning %d KafkaConsumers to %d partitions of Kafka topic %s for Subscription %s",
      totalConsumers, partitionInfo.size(), kafkaTopicName, subscription.getName());
  for (int i = 0; i < totalConsumers; i++) {
    Consumer<String, ByteBuffer> consumer =
        i == 0 ? first : kafkaClientFactory.createConsumer(subscription.getName());
    int consumerIndex = i;
    Set<TopicPartition> partitionSet =
        partitionInfo
            .stream()
            .filter(p -> p.partition() % totalConsumers == consumerIndex)
            .map(p -> new TopicPartition(kafkaTopicName, p.partition()))
            .collect(Collectors.toSet());
    consumer.assign(partitionSet);
    Map<TopicPartition, Long> endOffsets = consumer.endOffsets(partitionSet);
    for (TopicPartition tp : partitionSet) {
      committedOffsets.put(tp, consumer.committed(tp));
      logger.atFine().log(
          "%s assigned KafkaConsumer %d to %s:%d (End: %d, Committed %d)",
          subscription.getName(),
          consumerIndex,
          tp.topic(),
          tp.partition(),
          endOffsets.get(tp),
          Optional.ofNullable(consumer.committed(tp)).map(OffsetAndMetadata::offset).orElse(0L));
    }
    kafkaConsumers.add(new SynchronizedConsumer(consumer));
  }
}
 

public static void handleSeekOffset(ConsumerLocalConfigs effectiveLocal, Consumer consumer) {
    String seek = effectiveLocal.getSeek();
    if (!isEmpty(seek)) {
        String[] seekPosition = effectiveLocal.getSeekTopicPartitionOffset();
        TopicPartition topicPartition = new TopicPartition(seekPosition[0], parseInt(seekPosition[1]));

        Set<TopicPartition> topicPartitions = new HashSet<>();
        topicPartitions.add(topicPartition);

        consumer.unsubscribe();
        consumer.assign(topicPartitions);
        consumer.seek(topicPartition, parseLong(seekPosition[2]));
    }
}
 

public static Consumer<String, byte[]> createConsumer(Properties props, String subscribeTopic) {
    List<TopicPartition> topicPartitions = new ArrayList<>();
    for (String topic : subscribeTopic.split(",")) {
        topicPartitions.add(new TopicPartition(topic, 0));
    }
    Consumer<String, byte[]> consumer = new KafkaConsumer<>(props);
    consumer.assign(topicPartitions);
    consumer.seekToEnd(topicPartitions);
    return consumer;
}
 

public static Consumer<String, byte[]> createConsumer(Properties props, String subscribeTopic) throws Exception {
    TopicPartition topicPartition = new TopicPartition(subscribeTopic, 0);
    List<TopicPartition> topicPartitions = Arrays.asList(topicPartition);
    Consumer<String, byte[]> consumer = new KafkaConsumer<>(props);
    // consumer.subscribe(Arrays.asList(subscribeTopics.split(",")));
    consumer.assign(topicPartitions);
    // consumer都是在topo启动时创建。当Topo重启，目前的策略是对于kafka中未处理的msg放弃。不再消费。所以seek to end。
    consumer.seekToEnd(topicPartitions);
    return consumer;
}
 

/**
 * 取得尾部offset（LEO）
 *
 * @param topic
 * @param partition
 * @return
 */
private long findLogEndOffset(String topic, int partition) {
    Consumer consumer = getConsumer();
    TopicPartition topicPartition = new TopicPartition(topic, partition);
    List tpList = new ArrayList();
    tpList.add(topicPartition);
    consumer.assign(tpList);
    consumer.seekToEnd(tpList);
    Long longEndOffset = consumer.position(topicPartition);
    return longEndOffset;
}
 

public CommandStore(
    String commandTopic,
    Consumer<CommandId, Command> commandConsumer,
    Producer<CommandId, Command> commandProducer,
    CommandIdAssigner commandIdAssigner
) {
  this.commandTopic = commandTopic;
  this.commandConsumer = commandConsumer;
  this.commandProducer = commandProducer;
  this.commandIdAssigner = commandIdAssigner;

  commandConsumer.assign(Collections.singleton(new TopicPartition(commandTopic, 0)));

  closed = new AtomicBoolean(false);
}
 

@Test
public void doTest() throws Exception {
    System.out.println("==== KeycloakRefreshTokenWithJwtValidationTest ====");

    final String topic = "KeycloakRefreshTokenWithJwtValidationTest";
    final String tokenEndpointUri = "http://" + HOST + ":8080/auth/realms/" + REALM + "/protocol/openid-connect/token";

    String refreshToken = loginWithClientSecretForRefreshToken(URI.create(tokenEndpointUri), CLIENT_ID, CLIENT_SECRET);

    Properties defaults = new Properties();
    defaults.setProperty(ClientConfig.OAUTH_TOKEN_ENDPOINT_URI, tokenEndpointUri);
    defaults.setProperty(ClientConfig.OAUTH_CLIENT_ID, CLIENT_ID);
    defaults.setProperty(ClientConfig.OAUTH_CLIENT_SECRET, CLIENT_SECRET);
    defaults.setProperty(ClientConfig.OAUTH_REFRESH_TOKEN, refreshToken);
    defaults.setProperty(ClientConfig.OAUTH_USERNAME_CLAIM, "preferred_username");

    ConfigProperties.resolveAndExportToSystemProperties(defaults);

    Properties producerProps = buildProducerConfig();
    Producer<String, String> producer = new KafkaProducer<>(producerProps);

    producer.send(new ProducerRecord<>(topic, "The Message")).get();
    System.out.println("Produced The Message");

    Properties consumerProps = buildConsumerConfig();
    Consumer<String, String> consumer = new KafkaConsumer<>(consumerProps);

    TopicPartition partition = new TopicPartition(topic, 0);
    consumer.assign(Arrays.asList(partition));

    while (consumer.partitionsFor(topic, Duration.ofSeconds(1)).size() == 0) {
        System.out.println("No assignment yet for consumer");
    }
    consumer.seekToBeginning(Arrays.asList(partition));

    ConsumerRecords<String, String> records = consumer.poll(Duration.ofSeconds(1));

    Assert.assertEquals("Got message", 1, records.count());
    Assert.assertEquals("Is message text: 'The Message'", "The Message", records.iterator().next().value());
}
 

static void consume(Consumer<String, String> consumer, String topic) {
    TopicPartition partition = new TopicPartition(topic, 0);
    consumer.assign(Arrays.asList(partition));

    while (consumer.partitionsFor(topic, Duration.ofSeconds(1)).size() == 0) {
        System.out.println("No assignment yet for consumer");
    }

    consumer.seekToBeginning(Arrays.asList(partition));
    ConsumerRecords<String, String> records = consumer.poll(Duration.ofSeconds(10));

    Assert.assertTrue("Got message", records.count() >= 1);
}
 

@Test
public void doTest() throws Exception {
    System.out.println("==== KeycloakAccessTokenWithIntrospectValidationTest ====");

    final String topic = "KeycloakAccessTokenWithIntrospectValidationTest";
    final String tokenEndpointUri = "http://" + HOST + ":8080/auth/realms/" + REALM + "/protocol/openid-connect/token";

    // first, request access token using client id and secret
    TokenInfo info = OAuthAuthenticator.loginWithClientSecret(URI.create(tokenEndpointUri), null, null, CLIENT_ID, CLIENT_SECRET, true, null, null);


    Properties defaults = new Properties();
    defaults.setProperty(ClientConfig.OAUTH_TOKEN_ENDPOINT_URI, tokenEndpointUri);
    defaults.setProperty(ClientConfig.OAUTH_ACCESS_TOKEN, info.token());
    defaults.setProperty(ClientConfig.OAUTH_USERNAME_CLAIM, "preferred_username");

    ConfigProperties.resolveAndExportToSystemProperties(defaults);

    Properties producerProps = buildProducerConfig();
    Producer<String, String> producer = new KafkaProducer<>(producerProps);

    producer.send(new ProducerRecord<>(topic, "The Message")).get();
    System.out.println("Produced The Message");

    Properties consumerProps = buildConsumerConfig();
    Consumer<String, String> consumer = new KafkaConsumer<>(consumerProps);

    TopicPartition partition = new TopicPartition(topic, 0);
    consumer.assign(Arrays.asList(partition));

    while (consumer.partitionsFor(topic, Duration.ofSeconds(1)).size() == 0) {
        System.out.println("No assignment yet for consumer");
    }
    consumer.seekToBeginning(Arrays.asList(partition));

    ConsumerRecords<String, String> records = consumer.poll(Duration.ofSeconds(1));

    Assert.assertEquals("Got message", 1, records.count());
    Assert.assertEquals("Is message text: 'The Message'", "The Message", records.iterator().next().value());
}
 

public static void unassignConsumer(Consumer<String, byte[]> consumer) {
    consumer.assign(new ArrayList<>());
    consumer.close();
}
 

@Test
public void doTest() throws Exception {
    System.out.println("==== KeycloakClientCredentialsWithJwtValidationTest + test EC ====");

    Properties p = System.getProperties();
    for (Object key: p.keySet()) {
        System.out.println("" + key + "=" + p.get(key));
    }

    final String topic = "KeycloakClientCredentialsWithJwtValidationTest";
    final String tokenEndpointUri = "http://" + HOST + ":8080/auth/realms/" + REALM + "/protocol/openid-connect/token";

    Properties defaults = new Properties();
    defaults.setProperty(ClientConfig.OAUTH_TOKEN_ENDPOINT_URI, tokenEndpointUri);
    defaults.setProperty(ClientConfig.OAUTH_CLIENT_ID, "kafka-producer-client");
    defaults.setProperty(ClientConfig.OAUTH_CLIENT_SECRET, "kafka-producer-client-secret");
    defaults.setProperty(ClientConfig.OAUTH_USERNAME_CLAIM, "preferred_username");

    ConfigProperties.resolveAndExportToSystemProperties(defaults);

    Properties producerProps = buildProducerConfig();
    Producer<String, String> producer = new KafkaProducer<>(producerProps);

    producer.send(new ProducerRecord<>(topic, "The Message")).get();
    System.out.println("Produced The Message");

    Properties consumerProps = buildConsumerConfig();
    Consumer<String, String> consumer = new KafkaConsumer<>(consumerProps);

    TopicPartition partition = new TopicPartition(topic, 0);
    consumer.assign(Arrays.asList(partition));

    while (consumer.partitionsFor(topic, Duration.ofSeconds(1)).size() == 0) {
        System.out.println("No assignment yet for consumer");
    }
    consumer.seekToBeginning(Arrays.asList(partition));

    ConsumerRecords<String, String> records = consumer.poll(Duration.ofSeconds(1));

    Assert.assertEquals("Got message", 1, records.count());
    Assert.assertEquals("Is message text: 'The Message'", "The Message", records.iterator().next().value());
}
 

@Test
public void doTest() throws Exception {
    System.out.println("==== HydraOpaqueAccessTokenWithIntrospectValidationTest ====");

    final String topic = "HydraOpaqueAccessTokenWithIntrospectValidationTest";
    final String tokenEndpointUri = "https://" + HOST + ":4444/oauth2/token";

    Properties defaults = new Properties();
    defaults.setProperty(ClientConfig.OAUTH_TOKEN_ENDPOINT_URI, tokenEndpointUri);
    defaults.setProperty(ClientConfig.OAUTH_ACCESS_TOKEN_IS_JWT, "false");

    defaults.setProperty(ClientConfig.OAUTH_SSL_TRUSTSTORE_LOCATION, "../docker/target/kafka/certs/ca-truststore.p12");
    defaults.setProperty(ClientConfig.OAUTH_SSL_TRUSTSTORE_PASSWORD, "changeit");
    defaults.setProperty(ClientConfig.OAUTH_SSL_TRUSTSTORE_TYPE, "pkcs12");

    ConfigProperties.resolveAndExportToSystemProperties(defaults);

    // Request access token using client id and secret, and trustore configuration
    TokenInfo info = OAuthAuthenticator.loginWithClientSecret(URI.create(tokenEndpointUri),
            ConfigUtil.createSSLFactory(new ClientConfig()),
            null, CLIENT_ID, CLIENT_SECRET, true, null, null);

    // Configure received token for Kafka client auth
    defaults.setProperty(ClientConfig.OAUTH_ACCESS_TOKEN, info.token());
    ConfigProperties.resolveAndExportToSystemProperties(defaults);


    Properties producerProps = buildProducerConfig();
    Producer<String, String> producer = new KafkaProducer<>(producerProps);

    producer.send(new ProducerRecord<>(topic, "The Message")).get();
    System.out.println("Produced The Message");

    Properties consumerProps = buildConsumerConfig();
    Consumer<String, String> consumer = new KafkaConsumer<>(consumerProps);

    TopicPartition partition = new TopicPartition(topic, 0);
    consumer.assign(Arrays.asList(partition));

    while (consumer.partitionsFor(topic, Duration.ofSeconds(1)).size() == 0) {
        System.out.println("No assignment yet for consumer");
    }
    consumer.seekToBeginning(Arrays.asList(partition));

    ConsumerRecords<String, String> records = consumer.poll(Duration.ofSeconds(1));

    Assert.assertEquals("Got message", 1, records.count());
    Assert.assertEquals("Is message text: 'The Message'", "The Message", records.iterator().next().value());
}
 

@Test
public void doTest() throws Exception {
    System.out.println("==== HydraClientCredentialsWithJwtValidationTest ====");

    final String topic = "HydraClientCredentialsWithJwtValidationTest";
    final String tokenEndpointUri = "https://" + HOST + ":4444/oauth2/token";

    Properties defaults = new Properties();
    defaults.setProperty(ClientConfig.OAUTH_TOKEN_ENDPOINT_URI, tokenEndpointUri);
    defaults.setProperty(ClientConfig.OAUTH_CLIENT_ID, "kafka-producer-client");
    defaults.setProperty(ClientConfig.OAUTH_CLIENT_SECRET, "kafka-producer-client-secret");
    defaults.setProperty(ClientConfig.OAUTH_TOKEN_ENDPOINT_URI, tokenEndpointUri);

    defaults.setProperty(ClientConfig.OAUTH_SSL_TRUSTSTORE_LOCATION, "../docker/target/kafka/certs/ca-truststore.p12");
    defaults.setProperty(ClientConfig.OAUTH_SSL_TRUSTSTORE_PASSWORD, "changeit");
    defaults.setProperty(ClientConfig.OAUTH_SSL_TRUSTSTORE_TYPE, "pkcs12");

    ConfigProperties.resolveAndExportToSystemProperties(defaults);

    Properties producerProps = buildProducerConfig();
    Producer<String, String> producer = new KafkaProducer<>(producerProps);

    producer.send(new ProducerRecord<>(topic, "The Message")).get();
    System.out.println("Produced The Message");

    Properties consumerProps = buildConsumerConfig();
    Consumer<String, String> consumer = new KafkaConsumer<>(consumerProps);

    TopicPartition partition = new TopicPartition(topic, 0);
    consumer.assign(Arrays.asList(partition));

    while (consumer.partitionsFor(topic, Duration.ofSeconds(1)).size() == 0) {
        System.out.println("No assignment yet for consumer");
    }
    consumer.seekToBeginning(Arrays.asList(partition));

    ConsumerRecords<String, String> records = consumer.poll(Duration.ofSeconds(1));

    Assert.assertEquals("Got message", 1, records.count());
    Assert.assertEquals("Is message text: 'The Message'", "The Message", records.iterator().next().value());
}