Search in sources :

Example 1 with KafkaAckable

use of io.mantisrx.connector.kafka.KafkaAckable in project mantis by Netflix.

the class KafkaSourceTest method testKafkaSourceSingleConsumerReadsAllMessagesInOrderFromSinglePartition.

@Test
public void testKafkaSourceSingleConsumerReadsAllMessagesInOrderFromSinglePartition() throws InterruptedException {
    String testTopic = "testTopic" + topicNum.incrementAndGet();
    int numPartitions = 1;
    kafkaServer.createTopic(testTopic, numPartitions);
    int numMessages = 10;
    for (int i = 0; i < numMessages; i++) {
        ProducerRecord<String, String> keyedMessage = new ProducerRecord<>(testTopic, "{\"messageNum\":" + i + "}");
        kafkaServer.sendMessages(keyedMessage);
    }
    KafkaSource kafkaSource = new KafkaSource(new NoopRegistry());
    Context context = mock(Context.class);
    Parameters params = ParameterTestUtils.createParameters(KafkaSourceParameters.TOPIC, testTopic, KafkaSourceParameters.PREFIX + ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest", KafkaSourceParameters.PREFIX + ConsumerConfig.GROUP_ID_CONFIG, "testKafkaConsumer-" + random.nextInt());
    when(context.getParameters()).then((Answer<Parameters>) invocation -> params);
    when(context.getWorkerInfo()).then((Answer<WorkerInfo>) invocation -> new WorkerInfo("testJobName", "testJobName-1", 1, 0, 1, MantisJobDurationType.Perpetual, "1.1.1.1"));
    when(context.getJobId()).then((Answer<String>) invocation -> "testJobName-1");
    Index index = new Index(0, 10);
    Observable<Observable<KafkaAckable>> sourceObs = kafkaSource.call(context, index);
    final CountDownLatch latch = new CountDownLatch(numMessages);
    final AtomicInteger counter = new AtomicInteger(0);
    sourceObs.flatMap(kafkaAckableObs -> kafkaAckableObs).map(kafkaAckable -> {
        Optional<Map<String, Object>> parsedEvent = kafkaAckable.getKafkaData().getParsedEvent();
        assertTrue(parsedEvent.isPresent());
        assertEquals(counter.getAndIncrement(), parsedEvent.get().get("messageNum"));
        LOGGER.info("got message on topic {} consumer Id {}", parsedEvent.get(), kafkaAckable.getKafkaData().getMantisKafkaConsumerId());
        kafkaAckable.ack();
        latch.countDown();
        return parsedEvent;
    }).subscribe();
    assertTrue("timed out waiting to get all messages from Kafka", latch.await(10, TimeUnit.SECONDS));
    kafkaServer.deleteTopic(testTopic);
}
Also used : Context(io.mantisrx.runtime.Context) NoopRegistry(com.netflix.spectator.api.NoopRegistry) Index(io.mantisrx.runtime.source.Index) KafkaSourceParameters(io.mantisrx.connector.kafka.KafkaSourceParameters) ProducerRecord(org.apache.kafka.clients.producer.ProducerRecord) BeforeClass(org.junit.BeforeClass) MantisJobDurationType(io.mantisrx.runtime.MantisJobDurationType) LoggerFactory(org.slf4j.LoggerFactory) ParameterTestUtils(io.mantisrx.connector.kafka.ParameterTestUtils) Parameters(io.mantisrx.runtime.parameter.Parameters) Random(java.util.Random) Observable(rx.Observable) Answer(org.mockito.stubbing.Answer) AtomicInteger(java.util.concurrent.atomic.AtomicInteger) KafkaAckable(io.mantisrx.connector.kafka.KafkaAckable) Map(java.util.Map) AfterClass(org.junit.AfterClass) Logger(org.slf4j.Logger) TestCase.fail(junit.framework.TestCase.fail) ConcurrentHashMap(java.util.concurrent.ConcurrentHashMap) Assert.assertTrue(org.junit.Assert.assertTrue) Set(java.util.Set) ConsumerConfig(org.apache.kafka.clients.consumer.ConsumerConfig) Test(org.junit.Test) Mockito.when(org.mockito.Mockito.when) Context(io.mantisrx.runtime.Context) TimeUnit(java.util.concurrent.TimeUnit) CountDownLatch(java.util.concurrent.CountDownLatch) ConcurrentSkipListSet(java.util.concurrent.ConcurrentSkipListSet) Ignore(org.junit.Ignore) KafkaUnit(info.batey.kafka.unit.KafkaUnit) Optional(java.util.Optional) Assert.assertEquals(org.junit.Assert.assertEquals) Mockito.mock(org.mockito.Mockito.mock) WorkerInfo(io.mantisrx.runtime.WorkerInfo) KafkaSourceParameters(io.mantisrx.connector.kafka.KafkaSourceParameters) Parameters(io.mantisrx.runtime.parameter.Parameters) Optional(java.util.Optional) WorkerInfo(io.mantisrx.runtime.WorkerInfo) Index(io.mantisrx.runtime.source.Index) CountDownLatch(java.util.concurrent.CountDownLatch) Observable(rx.Observable) AtomicInteger(java.util.concurrent.atomic.AtomicInteger) NoopRegistry(com.netflix.spectator.api.NoopRegistry) ProducerRecord(org.apache.kafka.clients.producer.ProducerRecord) Test(org.junit.Test)

Example 2 with KafkaAckable

use of io.mantisrx.connector.kafka.KafkaAckable in project mantis by Netflix.

the class KafkaSource method createBackPressuredConsumerObs.

/**
 * Create an observable with back pressure semantics from the consumer records fetched using consumer.
 *
 * @param mantisKafkaConsumer non thread-safe KafkaConsumer
 * @param kafkaSourceConfig   configuration for the Mantis Kafka Source
 */
private Observable<KafkaAckable> createBackPressuredConsumerObs(final MantisKafkaConsumer<?> mantisKafkaConsumer, final MantisKafkaSourceConfig kafkaSourceConfig) {
    CheckpointStrategy checkpointStrategy = mantisKafkaConsumer.getStrategy();
    final CheckpointTrigger trigger = mantisKafkaConsumer.getTrigger();
    final ConsumerMetrics consumerMetrics = mantisKafkaConsumer.getConsumerMetrics();
    final TopicPartitionStateManager partitionStateManager = mantisKafkaConsumer.getPartitionStateManager();
    int mantisKafkaConsumerId = mantisKafkaConsumer.getConsumerId();
    SyncOnSubscribe<Iterator<ConsumerRecord<String, byte[]>>, KafkaAckable> syncOnSubscribe = SyncOnSubscribe.createStateful(() -> {
        final ConsumerRecords<String, byte[]> records = mantisKafkaConsumer.poll(kafkaSourceConfig.getConsumerPollTimeoutMs());
        if (LOGGER.isDebugEnabled()) {
            LOGGER.debug("topic listing for consumer {}", mantisKafkaConsumer.listTopics());
        }
        LOGGER.info("consumer subscribed to topic-partitions {}", mantisKafkaConsumer.assignment());
        return records.iterator();
    }, (consumerRecordIterator, observer) -> {
        Iterator<ConsumerRecord<String, byte[]>> it = consumerRecordIterator;
        final Set<TopicPartition> partitions = mantisKafkaConsumer.assignment();
        if (trigger.shouldCheckpoint()) {
            long startTime = System.currentTimeMillis();
            final Map<TopicPartition, OffsetAndMetadata> checkpoint = partitionStateManager.createCheckpoint(partitions);
            checkpointStrategy.persistCheckpoint(checkpoint);
            long now = System.currentTimeMillis();
            consumerMetrics.recordCheckpointDelay(now - startTime);
            consumerMetrics.incrementCommitCount();
            trigger.reset();
        }
        if (!done.get()) {
            try {
                if (!consumerRecordIterator.hasNext()) {
                    final ConsumerRecords<String, byte[]> consumerRecords = mantisKafkaConsumer.poll(kafkaSourceConfig.getConsumerPollTimeoutMs());
                    if (LOGGER.isDebugEnabled()) {
                        LOGGER.debug("poll returned {} records", consumerRecords.count());
                    }
                    it = consumerRecords.iterator();
                }
                if (it.hasNext()) {
                    final ConsumerRecord<String, byte[]> m = it.next();
                    final TopicPartition topicPartition = new TopicPartition(m.topic(), m.partition());
                    consumerMetrics.incrementInCount();
                    if (LOGGER.isDebugEnabled()) {
                        LOGGER.debug("updating read offset to " + m.offset() + " read " + m.value());
                    }
                    if (m.value() != null) {
                        try {
                            trigger.update(getPayloadSize(m));
                            if (kafkaSourceConfig.getParseMessageInSource()) {
                                final Parser parser = ParserType.parser(kafkaSourceConfig.getMessageParserType()).getParser();
                                if (parser.canParse(m.value())) {
                                    final Map<String, Object> parsedKafkaValue = parser.parseMessage(m.value());
                                    final KafkaData kafkaData = new KafkaData(m, Optional.ofNullable(parsedKafkaValue), Optional.ofNullable(m.key()), mantisKafkaConsumerId);
                                    final KafkaAckable ackable = new KafkaAckable(kafkaData, ackSubject);
                                    // record offset consumed in TopicPartitionStateManager before onNext to avoid race condition with Ack being processed before the consume is recorded
                                    partitionStateManager.recordMessageRead(topicPartition, m.offset());
                                    consumerMetrics.recordReadOffset(topicPartition, m.offset());
                                    observer.onNext(ackable);
                                } else {
                                    consumerMetrics.incrementParseFailureCount();
                                }
                            } else {
                                final KafkaData kafkaData = new KafkaData(m, Optional.empty(), Optional.ofNullable(m.key()), mantisKafkaConsumerId);
                                final KafkaAckable ackable = new KafkaAckable(kafkaData, ackSubject);
                                // record offset consumed in TopicPartitionStateManager before onNext to avoid race condition with Ack being processed before the consume is recorded
                                partitionStateManager.recordMessageRead(topicPartition, m.offset());
                                consumerMetrics.recordReadOffset(topicPartition, m.offset());
                                observer.onNext(ackable);
                            }
                        } catch (ParseException pe) {
                            consumerMetrics.incrementErrorCount();
                            LOGGER.warn("failed to parse {}:{} message {}", m.topic(), m.partition(), m.value(), pe);
                        }
                    } else {
                        consumerMetrics.incrementKafkaMessageValueNullCount();
                    }
                } else {
                    consumerMetrics.incrementWaitForDataCount();
                    if (LOGGER.isDebugEnabled()) {
                        LOGGER.debug("Reached head of partition, waiting for more data");
                    }
                    TimeUnit.MILLISECONDS.sleep(200);
                }
            } catch (TimeoutException toe) {
                consumerMetrics.incrementWaitForDataCount();
                if (LOGGER.isDebugEnabled()) {
                    LOGGER.debug("Reached head of partition waiting for more data");
                }
            } catch (OffsetOutOfRangeException oore) {
                LOGGER.warn("offsets out of range " + oore.partitions() + " will seek to beginning", oore);
                final Set<TopicPartition> topicPartitionSet = oore.partitions();
                for (TopicPartition tp : topicPartitionSet) {
                    LOGGER.info("partition {} consumer position {}", tp, mantisKafkaConsumer.position(tp));
                }
                mantisKafkaConsumer.seekToBeginning(oore.partitions().toArray(new TopicPartition[oore.partitions().size()]));
            } catch (InvalidRecordException ire) {
                consumerMetrics.incrementErrorCount();
                LOGGER.warn("iterator error with invalid message. message will be dropped " + ire.getMessage());
            } catch (KafkaException e) {
                consumerMetrics.incrementErrorCount();
                LOGGER.warn("Other Kafka exception, message will be dropped. " + e.getMessage());
            } catch (InterruptedException ie) {
                LOGGER.error("consumer interrupted", ie);
                Thread.currentThread().interrupt();
            } catch (Exception e) {
                consumerMetrics.incrementErrorCount();
                LOGGER.warn("caught exception", e);
            }
        } else {
            mantisKafkaConsumer.close();
        }
        return it;
    }, consumerRecordIterator -> {
        LOGGER.info("closing Kafka consumer on unsubscribe" + mantisKafkaConsumer.toString());
        mantisKafkaConsumer.close();
    });
    return Observable.create(syncOnSubscribe).subscribeOn(Schedulers.newThread()).doOnUnsubscribe(() -> LOGGER.info("consumer {} stopped due to unsubscribe", mantisKafkaConsumerId)).doOnError((t) -> {
        LOGGER.error("consumer {} stopped due to error", mantisKafkaConsumerId, t);
        consumerMetrics.incrementErrorCount();
    }).doOnTerminate(() -> LOGGER.info("consumer {} terminated", mantisKafkaConsumerId));
}
Also used : Index(io.mantisrx.runtime.source.Index) KafkaSourceParameters(io.mantisrx.connector.kafka.KafkaSourceParameters) org.apache.kafka.clients.consumer(org.apache.kafka.clients.consumer) java.util(java.util) CONSUMER_RECORD_OVERHEAD_BYTES(io.mantisrx.connector.kafka.source.MantisKafkaSourceConfig.CONSUMER_RECORD_OVERHEAD_BYTES) ParserType(io.mantisrx.connector.kafka.source.serde.ParserType) LoggerFactory(org.slf4j.LoggerFactory) AtomicBoolean(java.util.concurrent.atomic.AtomicBoolean) KafkaException(org.apache.kafka.common.KafkaException) Source(io.mantisrx.runtime.source.Source) SyncOnSubscribe(rx.observables.SyncOnSubscribe) InvalidRecordException(org.apache.kafka.common.record.InvalidRecordException) KafkaDataNotification(io.mantisrx.connector.kafka.KafkaDataNotification) CheckpointStrategyOptions(io.mantisrx.connector.kafka.source.checkpoint.strategy.CheckpointStrategyOptions) Observable(rx.Observable) BooleanParameter(io.mantisrx.runtime.parameter.type.BooleanParameter) KafkaData(io.mantisrx.connector.kafka.KafkaData) KafkaAckable(io.mantisrx.connector.kafka.KafkaAckable) SerializedSubject(rx.subjects.SerializedSubject) Schedulers(rx.schedulers.Schedulers) TopicPartition(org.apache.kafka.common.TopicPartition) TimeoutException(org.apache.kafka.common.errors.TimeoutException) Logger(org.slf4j.Logger) Parser(io.mantisrx.connector.kafka.source.serde.Parser) DEFAULT_PARSE_MSG_IN_SOURCE(io.mantisrx.connector.kafka.source.MantisKafkaSourceConfig.DEFAULT_PARSE_MSG_IN_SOURCE) CheckpointStrategy(io.mantisrx.connector.kafka.source.checkpoint.strategy.CheckpointStrategy) DEFAULT_NUM_KAFKA_CONSUMER_PER_WORKER(io.mantisrx.connector.kafka.source.MantisKafkaSourceConfig.DEFAULT_NUM_KAFKA_CONSUMER_PER_WORKER) ParseException(io.mantisrx.connector.kafka.source.serde.ParseException) Context(io.mantisrx.runtime.Context) DEFAULT_ENABLE_STATIC_PARTITION_ASSIGN(io.mantisrx.connector.kafka.source.MantisKafkaSourceConfig.DEFAULT_ENABLE_STATIC_PARTITION_ASSIGN) ParameterDefinition(io.mantisrx.runtime.parameter.ParameterDefinition) Validators(io.mantisrx.runtime.parameter.validator.Validators) TimeUnit(java.util.concurrent.TimeUnit) ConsumerMetrics(io.mantisrx.connector.kafka.source.metrics.ConsumerMetrics) DEFAULT_MAX_BYTES_IN_PROCESSING(io.mantisrx.connector.kafka.source.MantisKafkaSourceConfig.DEFAULT_MAX_BYTES_IN_PROCESSING) CheckpointTrigger(io.mantisrx.connector.kafka.source.checkpoint.trigger.CheckpointTrigger) Registry(com.netflix.spectator.api.Registry) IntParameter(io.mantisrx.runtime.parameter.type.IntParameter) StringParameter(io.mantisrx.runtime.parameter.type.StringParameter) PublishSubject(rx.subjects.PublishSubject) CheckpointStrategy(io.mantisrx.connector.kafka.source.checkpoint.strategy.CheckpointStrategy) InvalidRecordException(org.apache.kafka.common.record.InvalidRecordException) TimeoutException(org.apache.kafka.common.errors.TimeoutException) KafkaAckable(io.mantisrx.connector.kafka.KafkaAckable) KafkaException(org.apache.kafka.common.KafkaException) InvalidRecordException(org.apache.kafka.common.record.InvalidRecordException) TimeoutException(org.apache.kafka.common.errors.TimeoutException) ParseException(io.mantisrx.connector.kafka.source.serde.ParseException) Parser(io.mantisrx.connector.kafka.source.serde.Parser) TopicPartition(org.apache.kafka.common.TopicPartition) KafkaData(io.mantisrx.connector.kafka.KafkaData) KafkaException(org.apache.kafka.common.KafkaException) ParseException(io.mantisrx.connector.kafka.source.serde.ParseException) CheckpointTrigger(io.mantisrx.connector.kafka.source.checkpoint.trigger.CheckpointTrigger) ConsumerMetrics(io.mantisrx.connector.kafka.source.metrics.ConsumerMetrics)

Example 3 with KafkaAckable

use of io.mantisrx.connector.kafka.KafkaAckable in project mantis by Netflix.

the class KafkaSourceTest method testKafkaSourceMultipleConsumersReadsAllMessagesFromMultiplePartitions.

@Test
public void testKafkaSourceMultipleConsumersReadsAllMessagesFromMultiplePartitions() throws InterruptedException {
    String testTopic = "testTopic" + topicNum.incrementAndGet();
    int numPartitions = 2;
    kafkaServer.createTopic(testTopic, numPartitions);
    int numMessages = 10;
    Set<Integer> outstandingMsgs = new ConcurrentSkipListSet<>();
    for (int i = 0; i < numMessages; i++) {
        ProducerRecord<String, String> keyedMessage = new ProducerRecord<>(testTopic, "{\"messageNum\":" + i + "}");
        kafkaServer.sendMessages(keyedMessage);
        outstandingMsgs.add(i);
    }
    KafkaSource kafkaSource = new KafkaSource(new NoopRegistry());
    Context context = mock(Context.class);
    Parameters params = ParameterTestUtils.createParameters(KafkaSourceParameters.NUM_KAFKA_CONSUMER_PER_WORKER, 2, KafkaSourceParameters.TOPIC, testTopic, KafkaSourceParameters.PREFIX + ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest", KafkaSourceParameters.PREFIX + ConsumerConfig.GROUP_ID_CONFIG, "testKafkaConsumer-" + random.nextInt());
    when(context.getParameters()).then((Answer<Parameters>) invocation -> params);
    when(context.getWorkerInfo()).then((Answer<WorkerInfo>) invocation -> new WorkerInfo("testJobName", "testJobName-1", 1, 0, 1, MantisJobDurationType.Perpetual, "1.1.1.1"));
    when(context.getJobId()).then((Answer<String>) invocation -> "testJobName-1");
    Index index = new Index(0, 10);
    Observable<Observable<KafkaAckable>> sourceObs = kafkaSource.call(context, index);
    final CountDownLatch latch = new CountDownLatch(numMessages);
    final Map<Integer, Integer> lastMessageNumByConsumerId = new ConcurrentHashMap<>();
    sourceObs.flatMap(kafkaAckableObs -> kafkaAckableObs).map(kafkaAckable -> {
        Optional<Map<String, Object>> parsedEvent = kafkaAckable.getKafkaData().getParsedEvent();
        assertTrue(parsedEvent.isPresent());
        Integer messageNum = (Integer) parsedEvent.get().get("messageNum");
        assertTrue(outstandingMsgs.contains(messageNum));
        outstandingMsgs.remove(messageNum);
        int mantisKafkaConsumerId = kafkaAckable.getKafkaData().getMantisKafkaConsumerId();
        lastMessageNumByConsumerId.putIfAbsent(mantisKafkaConsumerId, -1);
        // assert consumption of higher message numbers across consumer instances
        assertTrue(messageNum > lastMessageNumByConsumerId.get(mantisKafkaConsumerId));
        lastMessageNumByConsumerId.put(mantisKafkaConsumerId, messageNum);
        LOGGER.info("got message on topic {} consumer id {}", parsedEvent.get(), mantisKafkaConsumerId);
        kafkaAckable.ack();
        latch.countDown();
        return parsedEvent;
    }).doOnError(t -> {
        LOGGER.error("caught unexpected exception", t);
        fail("test failed due to unexpected error " + t.getMessage());
    }).subscribe();
    assertTrue("timed out waiting to get all messages from Kafka", latch.await(10, TimeUnit.SECONDS));
    assertEquals(0, outstandingMsgs.size());
    assertTrue(lastMessageNumByConsumerId.keySet().size() == 2);
    lastMessageNumByConsumerId.keySet().forEach(consumerId -> {
        assertTrue(lastMessageNumByConsumerId.get(consumerId) >= 0);
    });
    kafkaServer.deleteTopic(testTopic);
}
Also used : Context(io.mantisrx.runtime.Context) NoopRegistry(com.netflix.spectator.api.NoopRegistry) Index(io.mantisrx.runtime.source.Index) KafkaSourceParameters(io.mantisrx.connector.kafka.KafkaSourceParameters) ProducerRecord(org.apache.kafka.clients.producer.ProducerRecord) BeforeClass(org.junit.BeforeClass) MantisJobDurationType(io.mantisrx.runtime.MantisJobDurationType) LoggerFactory(org.slf4j.LoggerFactory) ParameterTestUtils(io.mantisrx.connector.kafka.ParameterTestUtils) Parameters(io.mantisrx.runtime.parameter.Parameters) Random(java.util.Random) Observable(rx.Observable) Answer(org.mockito.stubbing.Answer) AtomicInteger(java.util.concurrent.atomic.AtomicInteger) KafkaAckable(io.mantisrx.connector.kafka.KafkaAckable) Map(java.util.Map) AfterClass(org.junit.AfterClass) Logger(org.slf4j.Logger) TestCase.fail(junit.framework.TestCase.fail) ConcurrentHashMap(java.util.concurrent.ConcurrentHashMap) Assert.assertTrue(org.junit.Assert.assertTrue) Set(java.util.Set) ConsumerConfig(org.apache.kafka.clients.consumer.ConsumerConfig) Test(org.junit.Test) Mockito.when(org.mockito.Mockito.when) Context(io.mantisrx.runtime.Context) TimeUnit(java.util.concurrent.TimeUnit) CountDownLatch(java.util.concurrent.CountDownLatch) ConcurrentSkipListSet(java.util.concurrent.ConcurrentSkipListSet) Ignore(org.junit.Ignore) KafkaUnit(info.batey.kafka.unit.KafkaUnit) Optional(java.util.Optional) Assert.assertEquals(org.junit.Assert.assertEquals) Mockito.mock(org.mockito.Mockito.mock) WorkerInfo(io.mantisrx.runtime.WorkerInfo) KafkaSourceParameters(io.mantisrx.connector.kafka.KafkaSourceParameters) Parameters(io.mantisrx.runtime.parameter.Parameters) ConcurrentSkipListSet(java.util.concurrent.ConcurrentSkipListSet) WorkerInfo(io.mantisrx.runtime.WorkerInfo) Index(io.mantisrx.runtime.source.Index) CountDownLatch(java.util.concurrent.CountDownLatch) Observable(rx.Observable) AtomicInteger(java.util.concurrent.atomic.AtomicInteger) NoopRegistry(com.netflix.spectator.api.NoopRegistry) ProducerRecord(org.apache.kafka.clients.producer.ProducerRecord) ConcurrentHashMap(java.util.concurrent.ConcurrentHashMap) Map(java.util.Map) ConcurrentHashMap(java.util.concurrent.ConcurrentHashMap) Test(org.junit.Test)

Example 4 with KafkaAckable

use of io.mantisrx.connector.kafka.KafkaAckable in project mantis by Netflix.

the class KafkaSourceTest method testKafkaSourceSingleConsumerHandlesMessageParseFailures.

@Test
public void testKafkaSourceSingleConsumerHandlesMessageParseFailures() throws InterruptedException {
    String testTopic = "testTopic" + topicNum.incrementAndGet();
    int numPartitions = 1;
    kafkaServer.createTopic(testTopic, numPartitions);
    int numMessages = 10;
    for (int i = 0; i < numMessages; i++) {
        ProducerRecord<String, String> keyedMessage = new ProducerRecord<>(testTopic, "{\"messageNum\":" + i + "}");
        kafkaServer.sendMessages(keyedMessage);
        ProducerRecord<String, String> invalidJsonMessage = new ProducerRecord<>(testTopic, "{\"messageNum:" + i + "}");
        kafkaServer.sendMessages(invalidJsonMessage);
    }
    KafkaSource kafkaSource = new KafkaSource(new NoopRegistry());
    Context context = mock(Context.class);
    Parameters params = ParameterTestUtils.createParameters(KafkaSourceParameters.TOPIC, testTopic, KafkaSourceParameters.PREFIX + ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest", KafkaSourceParameters.PREFIX + ConsumerConfig.GROUP_ID_CONFIG, "testKafkaConsumer-" + random.nextInt());
    when(context.getParameters()).then((Answer<Parameters>) invocation -> params);
    when(context.getWorkerInfo()).then((Answer<WorkerInfo>) invocation -> new WorkerInfo("testJobName", "testJobName-1", 1, 0, 1, MantisJobDurationType.Perpetual, "1.1.1.1"));
    when(context.getJobId()).then((Answer<String>) invocation -> "testJobName-1");
    Index index = new Index(0, 10);
    Observable<Observable<KafkaAckable>> sourceObs = kafkaSource.call(context, index);
    final CountDownLatch latch = new CountDownLatch(numMessages);
    final AtomicInteger counter = new AtomicInteger(0);
    sourceObs.flatMap(kafkaAckableObs -> kafkaAckableObs).map(kafkaAckable -> {
        Optional<Map<String, Object>> parsedEvent = kafkaAckable.getKafkaData().getParsedEvent();
        assertTrue(parsedEvent.isPresent());
        assertEquals(counter.getAndIncrement(), parsedEvent.get().get("messageNum"));
        LOGGER.info("got message on topic {} consumer Id {}", parsedEvent.get(), kafkaAckable.getKafkaData().getMantisKafkaConsumerId());
        kafkaAckable.ack();
        latch.countDown();
        return parsedEvent;
    }).subscribe();
    assertTrue("timed out waiting to get all messages from Kafka", latch.await(30, TimeUnit.SECONDS));
    kafkaServer.deleteTopic(testTopic);
}
Also used : Context(io.mantisrx.runtime.Context) NoopRegistry(com.netflix.spectator.api.NoopRegistry) Index(io.mantisrx.runtime.source.Index) KafkaSourceParameters(io.mantisrx.connector.kafka.KafkaSourceParameters) ProducerRecord(org.apache.kafka.clients.producer.ProducerRecord) BeforeClass(org.junit.BeforeClass) MantisJobDurationType(io.mantisrx.runtime.MantisJobDurationType) LoggerFactory(org.slf4j.LoggerFactory) ParameterTestUtils(io.mantisrx.connector.kafka.ParameterTestUtils) Parameters(io.mantisrx.runtime.parameter.Parameters) Random(java.util.Random) Observable(rx.Observable) Answer(org.mockito.stubbing.Answer) AtomicInteger(java.util.concurrent.atomic.AtomicInteger) KafkaAckable(io.mantisrx.connector.kafka.KafkaAckable) Map(java.util.Map) AfterClass(org.junit.AfterClass) Logger(org.slf4j.Logger) TestCase.fail(junit.framework.TestCase.fail) ConcurrentHashMap(java.util.concurrent.ConcurrentHashMap) Assert.assertTrue(org.junit.Assert.assertTrue) Set(java.util.Set) ConsumerConfig(org.apache.kafka.clients.consumer.ConsumerConfig) Test(org.junit.Test) Mockito.when(org.mockito.Mockito.when) Context(io.mantisrx.runtime.Context) TimeUnit(java.util.concurrent.TimeUnit) CountDownLatch(java.util.concurrent.CountDownLatch) ConcurrentSkipListSet(java.util.concurrent.ConcurrentSkipListSet) Ignore(org.junit.Ignore) KafkaUnit(info.batey.kafka.unit.KafkaUnit) Optional(java.util.Optional) Assert.assertEquals(org.junit.Assert.assertEquals) Mockito.mock(org.mockito.Mockito.mock) WorkerInfo(io.mantisrx.runtime.WorkerInfo) KafkaSourceParameters(io.mantisrx.connector.kafka.KafkaSourceParameters) Parameters(io.mantisrx.runtime.parameter.Parameters) Optional(java.util.Optional) WorkerInfo(io.mantisrx.runtime.WorkerInfo) Index(io.mantisrx.runtime.source.Index) CountDownLatch(java.util.concurrent.CountDownLatch) Observable(rx.Observable) AtomicInteger(java.util.concurrent.atomic.AtomicInteger) NoopRegistry(com.netflix.spectator.api.NoopRegistry) ProducerRecord(org.apache.kafka.clients.producer.ProducerRecord) Test(org.junit.Test)

Example 5 with KafkaAckable

use of io.mantisrx.connector.kafka.KafkaAckable in project mantis by Netflix.

the class KafkaSourceTest method testKafkaSourceMultipleConsumersStaticPartitionAssignment.

@Test
public void testKafkaSourceMultipleConsumersStaticPartitionAssignment() throws InterruptedException {
    String testTopic = "testTopic" + topicNum.incrementAndGet();
    int numConsumers = 3;
    int numPartitions = 3;
    kafkaServer.createTopic(testTopic, numPartitions);
    int numMessages = 10;
    Set<Integer> outstandingMsgs = new ConcurrentSkipListSet<>();
    for (int i = 0; i < numMessages; i++) {
        ProducerRecord<String, String> keyedMessage = new ProducerRecord<>(testTopic, "{\"messageNum\":" + i + "}");
        kafkaServer.sendMessages(keyedMessage);
        outstandingMsgs.add(i);
    }
    KafkaSource kafkaSource = new KafkaSource(new NoopRegistry());
    Context context = mock(Context.class);
    Parameters params = ParameterTestUtils.createParameters(KafkaSourceParameters.NUM_KAFKA_CONSUMER_PER_WORKER, numConsumers, KafkaSourceParameters.TOPIC, testTopic, KafkaSourceParameters.ENABLE_STATIC_PARTITION_ASSIGN, true, KafkaSourceParameters.TOPIC_PARTITION_COUNTS, testTopic + ":" + numPartitions, KafkaSourceParameters.PREFIX + ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest", KafkaSourceParameters.PREFIX + ConsumerConfig.GROUP_ID_CONFIG, "testKafkaConsumer-" + random.nextInt());
    when(context.getParameters()).then((Answer<Parameters>) invocation -> params);
    when(context.getWorkerInfo()).then((Answer<WorkerInfo>) invocation -> new WorkerInfo("testJobName", "testJobName-1", 1, 0, 1, MantisJobDurationType.Perpetual, "1.1.1.1"));
    when(context.getJobId()).then((Answer<String>) invocation -> "testJobName-1");
    // Force all consumer instances to be created on same JVM by setting total number of workers for this job to 1
    int totalNumWorkerForJob = 1;
    Index index = new Index(0, totalNumWorkerForJob);
    Observable<Observable<KafkaAckable>> sourceObs = kafkaSource.call(context, index);
    final CountDownLatch latch = new CountDownLatch(numMessages);
    final Map<Integer, Integer> lastMessageNumByConsumerId = new ConcurrentHashMap<>();
    sourceObs.flatMap(kafkaAckableObs -> kafkaAckableObs).map(kafkaAckable -> {
        Optional<Map<String, Object>> parsedEvent = kafkaAckable.getKafkaData().getParsedEvent();
        assertTrue(parsedEvent.isPresent());
        Integer messageNum = (Integer) parsedEvent.get().get("messageNum");
        assertTrue(outstandingMsgs.contains(messageNum));
        outstandingMsgs.remove(messageNum);
        int mantisKafkaConsumerId = kafkaAckable.getKafkaData().getMantisKafkaConsumerId();
        lastMessageNumByConsumerId.putIfAbsent(mantisKafkaConsumerId, -1);
        // assert consumption of higher message numbers across consumer instances
        assertTrue(messageNum > lastMessageNumByConsumerId.get(mantisKafkaConsumerId));
        lastMessageNumByConsumerId.put(mantisKafkaConsumerId, messageNum);
        LOGGER.info("got message on topic {} consumer id {}", parsedEvent.get(), mantisKafkaConsumerId);
        kafkaAckable.ack();
        latch.countDown();
        return parsedEvent;
    }).doOnError(t -> {
        LOGGER.error("caught unexpected exception", t);
        fail("test failed due to unexpected error " + t.getMessage());
    }).subscribe();
    assertTrue("timed out waiting to get all messages from Kafka", latch.await(10, TimeUnit.SECONDS));
    assertEquals(0, outstandingMsgs.size());
    assertTrue(lastMessageNumByConsumerId.keySet().size() == numConsumers);
    lastMessageNumByConsumerId.keySet().forEach(consumerId -> {
        assertTrue(lastMessageNumByConsumerId.get(consumerId) >= 0);
    });
    kafkaServer.deleteTopic(testTopic);
}
Also used : Context(io.mantisrx.runtime.Context) NoopRegistry(com.netflix.spectator.api.NoopRegistry) Index(io.mantisrx.runtime.source.Index) KafkaSourceParameters(io.mantisrx.connector.kafka.KafkaSourceParameters) ProducerRecord(org.apache.kafka.clients.producer.ProducerRecord) BeforeClass(org.junit.BeforeClass) MantisJobDurationType(io.mantisrx.runtime.MantisJobDurationType) LoggerFactory(org.slf4j.LoggerFactory) ParameterTestUtils(io.mantisrx.connector.kafka.ParameterTestUtils) Parameters(io.mantisrx.runtime.parameter.Parameters) Random(java.util.Random) Observable(rx.Observable) Answer(org.mockito.stubbing.Answer) AtomicInteger(java.util.concurrent.atomic.AtomicInteger) KafkaAckable(io.mantisrx.connector.kafka.KafkaAckable) Map(java.util.Map) AfterClass(org.junit.AfterClass) Logger(org.slf4j.Logger) TestCase.fail(junit.framework.TestCase.fail) ConcurrentHashMap(java.util.concurrent.ConcurrentHashMap) Assert.assertTrue(org.junit.Assert.assertTrue) Set(java.util.Set) ConsumerConfig(org.apache.kafka.clients.consumer.ConsumerConfig) Test(org.junit.Test) Mockito.when(org.mockito.Mockito.when) Context(io.mantisrx.runtime.Context) TimeUnit(java.util.concurrent.TimeUnit) CountDownLatch(java.util.concurrent.CountDownLatch) ConcurrentSkipListSet(java.util.concurrent.ConcurrentSkipListSet) Ignore(org.junit.Ignore) KafkaUnit(info.batey.kafka.unit.KafkaUnit) Optional(java.util.Optional) Assert.assertEquals(org.junit.Assert.assertEquals) Mockito.mock(org.mockito.Mockito.mock) WorkerInfo(io.mantisrx.runtime.WorkerInfo) KafkaSourceParameters(io.mantisrx.connector.kafka.KafkaSourceParameters) Parameters(io.mantisrx.runtime.parameter.Parameters) ConcurrentSkipListSet(java.util.concurrent.ConcurrentSkipListSet) WorkerInfo(io.mantisrx.runtime.WorkerInfo) Index(io.mantisrx.runtime.source.Index) CountDownLatch(java.util.concurrent.CountDownLatch) Observable(rx.Observable) AtomicInteger(java.util.concurrent.atomic.AtomicInteger) NoopRegistry(com.netflix.spectator.api.NoopRegistry) ProducerRecord(org.apache.kafka.clients.producer.ProducerRecord) ConcurrentHashMap(java.util.concurrent.ConcurrentHashMap) Map(java.util.Map) ConcurrentHashMap(java.util.concurrent.ConcurrentHashMap) Test(org.junit.Test)

Aggregations

KafkaAckable (io.mantisrx.connector.kafka.KafkaAckable)5 KafkaSourceParameters (io.mantisrx.connector.kafka.KafkaSourceParameters)5 Context (io.mantisrx.runtime.Context)5 Index (io.mantisrx.runtime.source.Index)5 TimeUnit (java.util.concurrent.TimeUnit)5 Logger (org.slf4j.Logger)5 LoggerFactory (org.slf4j.LoggerFactory)5 Observable (rx.Observable)5 NoopRegistry (com.netflix.spectator.api.NoopRegistry)4 KafkaUnit (info.batey.kafka.unit.KafkaUnit)4 ParameterTestUtils (io.mantisrx.connector.kafka.ParameterTestUtils)4 MantisJobDurationType (io.mantisrx.runtime.MantisJobDurationType)4 WorkerInfo (io.mantisrx.runtime.WorkerInfo)4 Parameters (io.mantisrx.runtime.parameter.Parameters)4 Map (java.util.Map)4 Optional (java.util.Optional)4 Random (java.util.Random)4 Set (java.util.Set)4 ConcurrentHashMap (java.util.concurrent.ConcurrentHashMap)4 ConcurrentSkipListSet (java.util.concurrent.ConcurrentSkipListSet)4