/** * computes partition for given record. * if the record has partition returns the value otherwise * calls configured partitioner class to compute the partition. */ private int partition(ProducerRecord<K, V> record, byte[] serializedKey, byte[] serializedValue, Cluster cluster) { Integer partition = record.partition(); return partition != null ? partition : partitioner.partition( record.topic(), record.key(), serializedKey, record.value(), serializedValue, cluster); }
/** * computes partition for given record. */ private int partition(ProducerRecord<K, V> record, Cluster cluster) { Integer partition = record.partition(); String topic = record.topic(); if (partition != null) { List<PartitionInfo> partitions = cluster.partitionsForTopic(topic); int numPartitions = partitions.size(); // they have given us a partition, use it if (partition < 0 || partition >= numPartitions) throw new IllegalArgumentException("Invalid partition given with record: " + partition + " is not in the range [0..." + numPartitions + "]."); return partition; } byte[] keyBytes = keySerializer.serialize(topic, record.headers(), record.key()); byte[] valueBytes = valueSerializer.serialize(topic, record.headers(), record.value()); return this.partitioner.partition(topic, record.key(), keyBytes, record.value(), valueBytes, cluster); }
@Test public void testRoundRobinWithUnavailablePartitions() { // When there are some unavailable partitions, we want to make sure that (1) we always pick an available partition, // and (2) the available partitions are selected in a round robin way. int countForPart0 = 0; int countForPart2 = 0; for (int i = 1; i <= 100; i++) { int part = partitioner.partition("test", null, null, null, null, cluster); assertTrue("We should never choose a leader-less node in round robin", part == 0 || part == 2); if (part == 0) countForPart0++; else countForPart2++; } assertEquals("The distribution between two available partitions should be even", countForPart0, countForPart2); }
@Test public void testPartitionSpread() throws Exception { Multiset<Integer> results = TreeMultiset.create(); Cluster c = Cluster.empty(); try (Partitioner p = new DefaultPartitioner()) { PartitionKeyGenerator pkg = new PartitionKeyGenerator(); mockPartitions(c); for (int i = 0; i < messages; i++) { results.add(p.partition("test", null, pkg.next(), null, null, c)); } int expected = messages / partitions; double threshold = expected * 0.05; for (Multiset.Entry<Integer> e : results.entrySet()) { int offBy = Math.abs(e.getCount() - expected); assertTrue("Partition " + e.getElement() + " had " + e.getCount() + " elements, expected " + expected + ", threshold is " + threshold, offBy < threshold); } } }