Examples with FunctionBasedQueueProducer org.apache.hudi.common.util.queue.FunctionBasedQueueProducer used on opensource projects

Example 1 with FunctionBasedQueueProducer

use of org.apache.hudi.common.util.queue.FunctionBasedQueueProducer in project hudi by apache.

the class TestBoundedInMemoryQueue method testCompositeProducerRecordReading.

/**
 * Test to ensure that we are reading all records from queue iterator when we have multiple producers.
 */
@SuppressWarnings("unchecked")
@Test
@Timeout(value = 60)
public void testCompositeProducerRecordReading() throws Exception {
    final int numRecords = 1000;
    final int numProducers = 40;
    final List<List<HoodieRecord>> recs = new ArrayList<>();
    final BoundedInMemoryQueue<HoodieRecord, HoodieLazyInsertIterable.HoodieInsertValueGenResult> queue = new BoundedInMemoryQueue(FileIOUtils.KB, getTransformFunction(HoodieTestDataGenerator.AVRO_SCHEMA));
    // Record Key to <Producer Index, Rec Index within a producer>
    Map<String, Tuple2<Integer, Integer>> keyToProducerAndIndexMap = new HashMap<>();
    for (int i = 0; i < numProducers; i++) {
        List<HoodieRecord> pRecs = dataGen.generateInserts(instantTime, numRecords);
        int j = 0;
        for (HoodieRecord r : pRecs) {
            assertFalse(keyToProducerAndIndexMap.containsKey(r.getRecordKey()));
            keyToProducerAndIndexMap.put(r.getRecordKey(), new Tuple2<>(i, j));
            j++;
        }
        recs.add(pRecs);
    }
    List<BoundedInMemoryQueueProducer<HoodieRecord>> producers = new ArrayList<>();
    for (int i = 0; i < recs.size(); i++) {
        final List<HoodieRecord> r = recs.get(i);
        // Alternate between pull and push based iterators
        if (i % 2 == 0) {
            producers.add(new IteratorBasedQueueProducer<>(r.iterator()));
        } else {
            producers.add(new FunctionBasedQueueProducer<>((buf) -> {
                Iterator<HoodieRecord> itr = r.iterator();
                while (itr.hasNext()) {
                    try {
                        buf.insertRecord(itr.next());
                    } catch (Exception e) {
                        throw new HoodieException(e);
                    }
                }
                return true;
            }));
        }
    }
    final List<Future<Boolean>> futureList = producers.stream().map(producer -> {
        return executorService.submit(() -> {
            producer.produce(queue);
            return true;
        });
    }).collect(Collectors.toList());
    // Close queue
    Future<Boolean> closeFuture = executorService.submit(() -> {
        try {
            for (Future f : futureList) {
                f.get();
            }
            queue.close();
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        return true;
    });
    // Used to ensure that consumer sees the records generated by a single producer in FIFO order
    Map<Integer, Integer> lastSeenMap = IntStream.range(0, numProducers).boxed().collect(Collectors.toMap(Function.identity(), x -> -1));
    Map<Integer, Integer> countMap = IntStream.range(0, numProducers).boxed().collect(Collectors.toMap(Function.identity(), x -> 0));
    // Read recs and ensure we have covered all producer recs.
    while (queue.iterator().hasNext()) {
        final HoodieLazyInsertIterable.HoodieInsertValueGenResult payload = queue.iterator().next();
        final HoodieRecord rec = payload.record;
        Tuple2<Integer, Integer> producerPos = keyToProducerAndIndexMap.get(rec.getRecordKey());
        Integer lastSeenPos = lastSeenMap.get(producerPos._1());
        countMap.put(producerPos._1(), countMap.get(producerPos._1()) + 1);
        lastSeenMap.put(producerPos._1(), lastSeenPos + 1);
        // Ensure we are seeing the next record generated
        assertEquals(lastSeenPos + 1, producerPos._2().intValue());
    }
    for (int i = 0; i < numProducers; i++) {
        // Ensure we have seen all the records for each producers
        assertEquals(Integer.valueOf(numRecords), countMap.get(i));
    }
    // Ensure Close future is done
    closeFuture.get();
}