/** * Create a new mock producer with invented metadata the given autoComplete setting and key\value serializers. * * Equivalent to {@link #MockProducer(Cluster, boolean, Partitioner, Serializer, Serializer)} new MockProducer(Cluster.empty(), autoComplete, new DefaultPartitioner(), keySerializer, valueSerializer)} */ public MockProducer(final boolean autoComplete, final Serializer<K> keySerializer, final Serializer<V> valueSerializer) { this(Cluster.empty(), autoComplete, new DefaultPartitioner(), keySerializer, valueSerializer); }
/** * Compute the partition for the given record. * * @param topic The topic name * @param key The key to partition on (or null if no key) * @param keyBytes serialized key to partition on (or null if no key) * @param value The value to partition on or null * @param valueBytes serialized value to partition on or null * @param cluster The current cluster metadata */ public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) { List<PartitionInfo> partitions = cluster.partitionsForTopic(topic); int numPartitions = partitions.size(); if (keyBytes == null) { int nextValue = nextValue(topic); List<PartitionInfo> availablePartitions = cluster.availablePartitionsForTopic(topic); if (availablePartitions.size() > 0) { int part = Utils.toPositive(nextValue) % availablePartitions.size(); return availablePartitions.get(part).partition(); } else { // no partitions are available, give a non-available partition return Utils.toPositive(nextValue) % numPartitions; } } else { // hash the keyBytes to choose a partition return Utils.toPositive(Utils.murmur2(keyBytes)) % numPartitions; } }
@Override public Integer partition(final String topic, final K key, final V value, final int numPartitions) { final byte[] keyBytes = keySerializer.serialize(topic, key); return defaultPartitioner.partition(topic, key, keyBytes, value, null, cluster); } }
@Test public void testPartitioner() throws Exception { PartitionInfo partitionInfo0 = new PartitionInfo(topic, 0, null, null, null); PartitionInfo partitionInfo1 = new PartitionInfo(topic, 1, null, null, null); Cluster cluster = new Cluster(null, new ArrayList<Node>(0), asList(partitionInfo0, partitionInfo1), Collections.<String>emptySet(), Collections.<String>emptySet()); MockProducer<String, String> producer = new MockProducer<>(cluster, true, new DefaultPartitioner(), new StringSerializer(), new StringSerializer()); ProducerRecord<String, String> record = new ProducerRecord<>(topic, "key", "value"); Future<RecordMetadata> metadata = producer.send(record); assertEquals("Partition should be correct", 1, metadata.get().partition()); producer.clear(); assertEquals("Clear should erase our history", 0, producer.history().size()); producer.close(); }
public DefaultStreamPartitioner(final Serializer<K> keySerializer, final Cluster cluster) { this.cluster = cluster; this.keySerializer = keySerializer; this.defaultPartitioner = new DefaultPartitioner(); }
@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); } } }