Java Code Examples for org.apache.kafka.clients.producer.ProducerRecord#partition()

@SuppressWarnings("unchecked")
@Override
public Object structuralValue(ProducerRecord<K, V> value) {
  if (consistentWithEquals()) {
    return value;
  } else {
    if (!ConsumerSpEL.hasHeaders()) {
      return new ProducerRecord<>(
          value.topic(), value.partition(), value.timestamp(), value.key(), value.value());
    } else {
      return new ProducerRecord<>(
          value.topic(),
          value.partition(),
          value.timestamp(),
          value.key(),
          value.value(),
          value.headers());
    }
  }
}
 

@Override
    public ProducerRecord<String, String> onSend(
            ProducerRecord<String, String> record) {
        String modifiedValue = "prefix1-" + record.value();
        return new ProducerRecord<>(record.topic(),
                record.partition(), record.timestamp(),
                record.key(), modifiedValue, record.headers());
//        if (record.value().length() < 5) {
//            throw new RuntimeException();
//        }
//        return record;
    }
 

@Override
public void onSendEnd(Span span, ProducerRecord producerRecord, KafkaProducer kafkaProducer, @Nullable Throwable throwable) {

    // Topic address collection is normally very fast, as it uses cached cluster state information. However,
    // when the cluster metadata is required to be updated, its query may block for a short period. In
    // addition, the discovery operation allocates two objects. Therefore, we have the ability to turn it off.
    if (messagingConfiguration.shouldCollectQueueAddress()) {
        try {
            // Better get the destination now, as if the partition's leader was replaced, it may be reflected at
            // this point.
            @SuppressWarnings("unchecked") List<PartitionInfo> partitions = kafkaProducer.partitionsFor(producerRecord.topic());
            Integer partition = producerRecord.partition();
            PartitionInfo partitionInfo = null;
            if (partition != null) {
                partitionInfo = partitions.get(partition);
            } else if (!partitions.isEmpty()) {
                // probably not a partitioned topic, so look for the singe entry
                partitionInfo = partitions.get(0);
            }
            if (partitionInfo != null) {
                // Records are always sent to the leader of a partition and then synced with replicas internally by
                // the broker.
                Node leader = partitionInfo.leader();
                if (leader != null) {
                    span.getContext().getDestination().withAddress(leader.host()).withPort(leader.port());
                }
            }
        } catch (Exception e) {
            logger.error("Failed to get Kafka producer's destination", e);
        }
    }

    span.captureException(throwable);

    // Not ending here- ending in the callback
    span.deactivate();
}
 

/**
 * 在将消息序列化和计算分区之前会调用
 *
 * @param record
 * @return
 */
@Override
public ProducerRecord<String, String> onSend(ProducerRecord<String, String> record) {
    // 拦截器，仅在 value 之前添加前缀
    String modifiedValue = "prefix1-" + record.value();
    return new ProducerRecord<>(record.topic(), record.partition(), record.timestamp(),
            record.key(), modifiedValue, record.headers());
}
 

public <K, V> void onSend(ProducerRecord<K, V> record, Span span) {
  setCommonTags(span);
  Tags.MESSAGE_BUS_DESTINATION.set(span, record.topic());
  if (record.partition() != null) {
    span.setTag("partition", record.partition());
  }
}
 

public static ProducerRecord prepareRecordToSend(String topicName, ProducerRecord recordToSend) {

        return new ProducerRecord(topicName,
                recordToSend.partition(),
                recordToSend.timestamp(),
                recordToSend.key(),
                recordToSend.value(),
                recordToSend.headers());
    }
 

@Test
public void testWriteLogMessagesWithCrcPartitioning() throws Exception {
  KafkaMessagePartitioner partitioner = new Crc32ByteArrayPartitioner();
  KafkaProducerConfig config = new KafkaProducerConfig();
  SingerSettings.setSingerConfig(new SingerConfig());
  KafkaProducerManager.injectTestProducer(config, producer);
  // default value for skip noleader partition is false
  KafkaWriter writer = new KafkaWriter(config, partitioner, "topicx", false, Executors.newCachedThreadPool());

  List<PartitionInfo> partitions = ImmutableList.copyOf(Arrays.asList(
      new PartitionInfo("topicx", 1, new Node(2, "broker2", 9092, "us-east-1b"), null, null),
      new PartitionInfo("topicx", 0, new Node(1, "broker1", 9092, "us-east-1a"), null, null),
      new PartitionInfo("topicx", 2, new Node(3, "broker3", 9092, "us-east-1c"), null, null),
      new PartitionInfo("topicx", 6, new Node(2, "broker2", 9092, "us-east-1b"), null, null),
      new PartitionInfo("topicx", 3, new Node(4, "broker4", 9092, "us-east-1a"), null, null),
      new PartitionInfo("topicx", 5, new Node(1, "broker1", 9092, "us-east-1a"), null, null),
      new PartitionInfo("topicx", 7, new Node(3, "broker3", 9092, "us-east-1c"), null, null),
      new PartitionInfo("topicx", 4, new Node(5, "broker5", 9092, "us-east-1b"), null, null),
      new PartitionInfo("topicx", 8, new Node(4, "broker4", 9092, "us-east-1a"), null, null),
      new PartitionInfo("topicx", 9, new Node(5, "broker5", 9092, "us-east-1b"), null, null),
      new PartitionInfo("topicx", 10, new Node(1, "broker1", 9092, "us-east-1a"), null, null)));

  when(producer.partitionsFor("topicx")).thenReturn(partitions);

  // message with same key will be put together in the same bucket (same partition);
  List<String> keys = IntStream.range(0, NUM_KEYS).mapToObj(i->"key"+i).collect(Collectors.toList());
  Map<Integer, List<LogMessage>> msgPartitionMap = new HashMap<>();
  Map<Integer, List<ProducerRecord<byte[], byte[]>>> recordPartitionMap = new HashMap<>();
  Map<Integer, List<RecordMetadata>> metadataPartitionMap = new HashMap<>();
  HashFunction crc32 = Hashing.crc32();
  List<LogMessage> logMessages = new ArrayList<>();
  for(int i = 0; i < NUM_KEYS; i++){
    for(int j = 0; j < NUM_EVENTS / NUM_KEYS; j++){
       LogMessage logMessage = new LogMessage();
       logMessage.setKey(keys.get(i).getBytes());
       logMessage.setMessage(ByteBuffer.allocate(100).put(String.valueOf(i).getBytes()));
       logMessages.add(logMessage);
       int partitionId = Math.abs(crc32.hashBytes(logMessage.getKey()).asInt() % partitions.size());
       ProducerRecord<byte[], byte[]> record = new ProducerRecord<byte[], byte[]>("topicx", partitionId, logMessage.getKey(), logMessage.getMessage());
       RecordMetadata recordMetadata = new RecordMetadata(new TopicPartition(  record.topic(),
          record.partition()), 0, 0, 0, 0L, record.key().length, record.value().length);
       when(producer.send(record)).thenReturn(ConcurrentUtils.constantFuture(recordMetadata));

       if (msgPartitionMap.containsKey(partitionId)){
         msgPartitionMap.get(partitionId).add(logMessage);
         recordPartitionMap.get(partitionId).add(record);
         metadataPartitionMap.get(partitionId).add(recordMetadata);
       } else {
         msgPartitionMap.put(partitionId, new ArrayList<>());
         recordPartitionMap.put(partitionId, new ArrayList<>());
         metadataPartitionMap.put(partitionId, new ArrayList<>());
         msgPartitionMap.get(partitionId).add(logMessage);
         recordPartitionMap.get(partitionId).add(record);
         metadataPartitionMap.get(partitionId).add(recordMetadata);
       }
    }
  }

  List<PartitionInfo> sortedPartitions = new ArrayList<>(partitions);
  Collections.sort(sortedPartitions, new PartitionComparator());

  Map<Integer, Map<Integer, LoggingAuditHeaders>> mapOfHeadersMap = new HashMap<>();
  Map<Integer, List<ProducerRecord<byte[], byte[]>>> messageCollation = writer.messageCollation(partitions, "topicx", logMessages, mapOfHeadersMap);

  for(int partitionId = 0; partitionId < messageCollation.keySet().size(); partitionId++) {
    if (messageCollation.get(partitionId).size() == 0) {
      continue;
    }
    List<ProducerRecord<byte[], byte[]>> writerOutput = messageCollation.get(partitionId);

    // verify the message order is what is expected by calling messageCollation()
    List<ProducerRecord<byte[], byte[]>> expectedRecords = recordPartitionMap.get(partitionId);
    assertEquals(expectedRecords.size(), writerOutput.size());
    for(int j = 0; j < writerOutput.size(); j++){
      assertEquals(expectedRecords.get(j), writerOutput.get(j));
    }

    // verify the content of LogMessage and the content of ProducerRecord match
    List<LogMessage> originalData = msgPartitionMap.get(partitionId);
    assertEquals(originalData.size(), writerOutput.size());
    for (int j = 0; j < writerOutput.size(); j++) {
      assertTrue(Arrays.equals(originalData.get(j).getKey(), writerOutput.get(j).key()));
      assertTrue(Arrays.equals(originalData.get(j).getMessage(), writerOutput.get(j).value()));
    }

    // verify the RecordMetadata that corresponds to record send to certain partitions are put
    // together into a list and the order of the RecordMetadata is same as the original message order
    List<RecordMetadata> expectedRecordMetadata = metadataPartitionMap.get(partitionId);
    KafkaWritingTaskResult kafkaWritingTaskResult  = writer.getClusterThreadPool().submit(new
        KafkaWritingTask(producer, writerOutput, 0, sortedPartitions)).get();
    assertEquals(expectedRecordMetadata.size(), kafkaWritingTaskResult.getRecordMetadataList().size());
    for(int j = 0; j < expectedRecordMetadata.size(); j++){
      assertEquals(expectedRecordMetadata.get(j), kafkaWritingTaskResult.getRecordMetadataList().get(j));
    }
  }

  // validate if writes are throwing any error
  writer.writeLogMessages(logMessages);
  writer.close();
}
 

private RecordMetadata getRecordMetadata(ProducerRecord record) {
  return new RecordMetadata(new TopicPartition(record.topic(), record.partition() == null ? 0 : record.partition()), 0, this.msgsSent.get(), -1L, -1L, -1, -1);
}
 

public FutureSuccess(ProducerRecord record, int offset) {
  this.record = record;
  this.metadata = new RecordMetadata(new TopicPartition(record.topic(), record.partition() == null ? 0 : record.partition()), 0, offset, RecordBatch.NO_TIMESTAMP, -1L, -1, -1);
}