Examples with Partitioner org.apache.kafka.clients.producer.Partitioner used on opensource projects

Example 1 with Partitioner

use of org.apache.kafka.clients.producer.Partitioner in project kafka by apache.

the class RecordAccumulatorTest method testStickyBatches.

@Test
public void testStickyBatches() throws Exception {
    long now = time.milliseconds();
    // Test case assumes that the records do not fill the batch completely
    int batchSize = 1025;
    Partitioner partitioner = new DefaultPartitioner();
    RecordAccumulator accum = createTestRecordAccumulator(3200, batchSize + DefaultRecordBatch.RECORD_BATCH_OVERHEAD, 10L * batchSize, CompressionType.NONE, 10);
    int expectedAppends = expectedNumAppendsNoKey(batchSize);
    // Create first batch
    int partition = partitioner.partition(topic, null, null, "value", value, cluster);
    TopicPartition tp = new TopicPartition(topic, partition);
    accum.append(tp, 0L, null, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds());
    int appends = 1;
    boolean switchPartition = false;
    while (!switchPartition) {
        // Append to the first batch
        partition = partitioner.partition(topic, null, null, "value", value, cluster);
        tp = new TopicPartition(topic, partition);
        RecordAccumulator.RecordAppendResult result = accum.append(tp, 0L, null, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, true, time.milliseconds());
        Deque<ProducerBatch> partitionBatches1 = accum.batches().get(tp1);
        Deque<ProducerBatch> partitionBatches2 = accum.batches().get(tp2);
        Deque<ProducerBatch> partitionBatches3 = accum.batches().get(tp3);
        int numBatches = (partitionBatches1 == null ? 0 : partitionBatches1.size()) + (partitionBatches2 == null ? 0 : partitionBatches2.size()) + (partitionBatches3 == null ? 0 : partitionBatches3.size());
        // Only one batch is created because the partition is sticky.
        assertEquals(1, numBatches);
        switchPartition = result.abortForNewBatch;
        // We only appended if we do not retry.
        if (!switchPartition) {
            appends++;
            assertEquals(0, accum.ready(cluster, now).readyNodes.size(), "No partitions should be ready.");
        }
    }
    // Batch should be full.
    assertEquals(1, accum.ready(cluster, time.milliseconds()).readyNodes.size());
    assertEquals(appends, expectedAppends);
    switchPartition = false;
    // KafkaProducer would call this method in this case, make second batch
    partitioner.onNewBatch(topic, cluster, partition);
    partition = partitioner.partition(topic, null, null, "value", value, cluster);
    tp = new TopicPartition(topic, partition);
    accum.append(tp, 0L, null, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, false, time.milliseconds());
    appends++;
    // These appends all go into the second batch
    while (!switchPartition) {
        partition = partitioner.partition(topic, null, null, "value", value, cluster);
        tp = new TopicPartition(topic, partition);
        RecordAccumulator.RecordAppendResult result = accum.append(tp, 0L, null, value, Record.EMPTY_HEADERS, null, maxBlockTimeMs, true, time.milliseconds());
        Deque<ProducerBatch> partitionBatches1 = accum.batches().get(tp1);
        Deque<ProducerBatch> partitionBatches2 = accum.batches().get(tp2);
        Deque<ProducerBatch> partitionBatches3 = accum.batches().get(tp3);
        int numBatches = (partitionBatches1 == null ? 0 : partitionBatches1.size()) + (partitionBatches2 == null ? 0 : partitionBatches2.size()) + (partitionBatches3 == null ? 0 : partitionBatches3.size());
        // Only two batches because the new partition is also sticky.
        assertEquals(2, numBatches);
        switchPartition = result.abortForNewBatch;
        // We only appended if we do not retry.
        if (!switchPartition) {
            appends++;
        }
    }
    // There should be two full batches now.
    assertEquals(appends, 2 * expectedAppends);
}

Example 2 with Partitioner

use of org.apache.kafka.clients.producer.Partitioner in project kafka by apache.

the class DefaultPartitionerTest method testKeyPartitionIsStable.

@Test
public void testKeyPartitionIsStable() {
    final Partitioner partitioner = new DefaultPartitioner();
    final Cluster cluster = new Cluster("clusterId", asList(NODES), PARTITIONS, Collections.<String>emptySet(), Collections.<String>emptySet());
    int partition = partitioner.partition("test", null, KEY_BYTES, null, null, cluster);
    assertEquals(partition, partitioner.partition("test", null, KEY_BYTES, null, null, cluster), "Same key should yield same partition");
}