Search in sources :

Example 1 with IRecordProcessorFactory

use of com.amazonaws.services.kinesis.clientlibrary.interfaces.v2.IRecordProcessorFactory in project samza by apache.

the class TestKinesisSystemConsumer method testProcessRecordsHelper.

/**
 * Helper to simulate and test the life-cycle of record processing from a kinesis stream with a given number of shards
 * 1. Creation of record processors.
 * 2. Initialization of record processors.
 * 3. Processing records via record processors.
 * 4. Calling checkpoint on record processors.
 * 5. Shutting down (due to re-assignment or lease expiration) record processors.
 */
private void testProcessRecordsHelper(String system, String stream, int numShards, int numRecordsPerShard) throws InterruptedException, NoSuchFieldException, IllegalAccessException {
    KinesisConfig kConfig = new KinesisConfig(new MapConfig());
    // Create consumer
    KinesisSystemConsumer consumer = new KinesisSystemConsumer(system, kConfig, new NoOpMetricsRegistry());
    initializeMetrics(consumer, stream);
    List<SystemStreamPartition> ssps = new LinkedList<>();
    IntStream.range(0, numShards).forEach(p -> {
        SystemStreamPartition ssp = new SystemStreamPartition(system, stream, new Partition(p));
        ssps.add(ssp);
    });
    ssps.forEach(ssp -> consumer.register(ssp, SYSTEM_CONSUMER_REGISTER_OFFSET));
    // Create Kinesis record processor factory
    IRecordProcessorFactory factory = consumer.createRecordProcessorFactory(stream);
    // Create and initialize Kinesis record processor
    Map<String, KinesisRecordProcessor> processorMap = createAndInitProcessors(factory, numShards);
    List<KinesisRecordProcessor> processorList = new ArrayList<>(processorMap.values());
    // Generate records to Kinesis record processor
    Map<KinesisRecordProcessor, List<Record>> inputRecordMap = generateRecords(numRecordsPerShard, processorList);
    // Verification steps
    // Read events from the BEM queue
    Map<SystemStreamPartition, List<IncomingMessageEnvelope>> messages = readEvents(new HashSet<>(ssps), consumer, numRecordsPerShard);
    if (numRecordsPerShard > 0) {
        Assert.assertEquals(messages.size(), numShards);
    } else {
        // No input records and hence no messages
        Assert.assertEquals(messages.size(), 0);
        return;
    }
    Map<SystemStreamPartition, KinesisRecordProcessor> sspToProcessorMap = getProcessorMap(consumer);
    ssps.forEach(ssp -> {
        try {
            KinesisRecordProcessor processor = sspToProcessorMap.get(ssp);
            // Verify that the read messages are received in order and are the same as input records
            Assert.assertEquals(messages.get(ssp).size(), numRecordsPerShard);
            List<IncomingMessageEnvelope> envelopes = messages.get(ssp);
            List<Record> inputRecords = inputRecordMap.get(processor);
            verifyRecords(envelopes, inputRecords, processor.getShardId());
            // Call checkpoint on consumer and verify that the checkpoint is called with the right offset
            IncomingMessageEnvelope lastEnvelope = envelopes.get(envelopes.size() - 1);
            consumer.afterCheckpoint(Collections.singletonMap(ssp, lastEnvelope.getOffset()));
            ArgumentCaptor<String> argument = ArgumentCaptor.forClass(String.class);
            verify(getCheckpointer(processor)).checkpoint(argument.capture());
            Assert.assertEquals(inputRecords.get(inputRecords.size() - 1).getSequenceNumber(), argument.getValue());
            // Call shutdown (with ZOMBIE reason) on processor and verify if shutdown freed the ssp mapping
            shutDownProcessor(processor, ShutdownReason.ZOMBIE);
            Assert.assertFalse(sspToProcessorMap.containsValue(processor));
            Assert.assertTrue(isSspAvailable(consumer, ssp));
        } catch (NoSuchFieldException | IllegalAccessException | InvalidStateException | ShutdownException ex) {
            throw new RuntimeException(ex);
        }
    });
}
Also used : KinesisConfig(org.apache.samza.system.kinesis.KinesisConfig) TestKinesisRecordProcessor(org.apache.samza.system.kinesis.consumer.TestKinesisRecordProcessor) IncomingMessageEnvelope(org.apache.samza.system.IncomingMessageEnvelope) ArrayList(java.util.ArrayList) ShutdownException(com.amazonaws.services.kinesis.clientlibrary.exceptions.ShutdownException) IRecordProcessorFactory(com.amazonaws.services.kinesis.clientlibrary.interfaces.v2.IRecordProcessorFactory) NoOpMetricsRegistry(org.apache.samza.util.NoOpMetricsRegistry) ArrayList(java.util.ArrayList) LinkedList(java.util.LinkedList) List(java.util.List) Record(com.amazonaws.services.kinesis.model.Record) MapConfig(org.apache.samza.config.MapConfig) SystemStreamPartition(org.apache.samza.system.SystemStreamPartition) Partition(org.apache.samza.Partition) InvalidStateException(com.amazonaws.services.kinesis.clientlibrary.exceptions.InvalidStateException) LinkedList(java.util.LinkedList) SystemStreamPartition(org.apache.samza.system.SystemStreamPartition)

Aggregations

InvalidStateException (com.amazonaws.services.kinesis.clientlibrary.exceptions.InvalidStateException)1 ShutdownException (com.amazonaws.services.kinesis.clientlibrary.exceptions.ShutdownException)1 IRecordProcessorFactory (com.amazonaws.services.kinesis.clientlibrary.interfaces.v2.IRecordProcessorFactory)1 Record (com.amazonaws.services.kinesis.model.Record)1 ArrayList (java.util.ArrayList)1 LinkedList (java.util.LinkedList)1 List (java.util.List)1 Partition (org.apache.samza.Partition)1 MapConfig (org.apache.samza.config.MapConfig)1 IncomingMessageEnvelope (org.apache.samza.system.IncomingMessageEnvelope)1 SystemStreamPartition (org.apache.samza.system.SystemStreamPartition)1 KinesisConfig (org.apache.samza.system.kinesis.KinesisConfig)1 TestKinesisRecordProcessor (org.apache.samza.system.kinesis.consumer.TestKinesisRecordProcessor)1 NoOpMetricsRegistry (org.apache.samza.util.NoOpMetricsRegistry)1