Examples with Index io.mantisrx.runtime.source.Index used on opensource projects

Example 11 with Index

use of io.mantisrx.runtime.source.Index in project mantis by Netflix.

the class StageExecutors method executeSource.

@SuppressWarnings({ "rawtypes", "unchecked" })
public static void executeSource(final int workerIndex, final SourceHolder source, final StageConfig stage, WorkerPublisher publisher, final Context context, final Observable<Integer> totalWorkerAtStageObservable) {
    // create a consumer from passed in source
    WorkerConsumer sourceConsumer = new WorkerConsumer() {

        @Override
        public Observable start(StageConfig stage) {
            Index index = new Index(workerIndex, totalWorkerAtStageObservable);
            // call init on source
            source.getSourceFunction().init(context, index);
            Observable<Observable<?>> sourceObservable = (Observable) source.getSourceFunction().call(context, new Index(workerIndex, totalWorkerAtStageObservable));
            return MantisMarker.sourceOut(sourceObservable);
        }

        @Override
        public void stop() {
        }
    };
    executeIntermediate(sourceConsumer, stage, publisher, context);
}

Example 12 with Index

use of io.mantisrx.runtime.source.Index in project mantis by Netflix.

the class StageExecutors method executeSingleStageJob.

@SuppressWarnings({ "rawtypes", "unchecked" })
public static void executeSingleStageJob(final SourceHolder source, final StageConfig stage, final SinkHolder sink, final PortSelector portSelector, RxMetrics rxMetrics, final Context context, Action0 sinkObservableTerminatedCallback, final int workerIndex, final Observable<Integer> totalWorkerAtStageObservable, final Action0 onSinkSubscribe, final Action0 onSinkUnsubscribe, Action0 observableOnCompleteCallback, Action1<Throwable> observableOnErrorCallback) {
    // no previous stage for single stage job
    // source consumer
    WorkerConsumer sourceConsumer = new WorkerConsumer() {

        @Override
        public Observable start(StageConfig previousStage) {
            Index index = new Index(workerIndex, totalWorkerAtStageObservable);
            // call init on source
            source.getSourceFunction().init(context, index);
            Observable<Observable<?>> sourceObservable = (Observable) source.getSourceFunction().call(context, index);
            if (stage.getInputStrategy() == StageConfig.INPUT_STRATEGY.CONCURRENT) {
                return sourceObservable;
            } else {
                return Observable.just(Observable.merge(sourceObservable));
            }
        }

        @Override
        public void stop() {
        }
    };
    // sink publisher with metrics
    WorkerPublisher sinkPublisher = new SinkPublisher(sink, portSelector, context, sinkObservableTerminatedCallback, onSinkSubscribe, onSinkUnsubscribe, observableOnCompleteCallback, observableOnErrorCallback);
    StageExecutors.executeIntermediate(sourceConsumer, stage, sinkPublisher, context);
}

Example 13 with Index

use of io.mantisrx.runtime.source.Index 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);
}

Example 14 with Index

use of io.mantisrx.runtime.source.Index 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);
}

Example 15 with Index

use of io.mantisrx.runtime.source.Index 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);
}