Examples with WatermarkEstimatorFactory cz.o2.proxima.time.WatermarkEstimatorFactory used on opensource projects

Example 1 with WatermarkEstimatorFactory

use of cz.o2.proxima.time.WatermarkEstimatorFactory in project proxima-platform by O2-Czech-Republic.

the class KafkaLogReader method listener.

// create rebalance listener from consumer
private ConsumerRebalanceListener listener(String name, AtomicReference<KafkaConsumer<Object, Object>> kafka, ElementConsumer<Object, Object> consumer, Map<TopicPartition, Integer> emptyPollCount, Map<TopicPartition, Integer> topicPartitionToId, AtomicReference<PartitionedWatermarkEstimator> watermarkEstimator) {
    return new ConsumerRebalanceListener() {

        private final Set<TopicPartition> currentlyAssigned = new HashSet<>();

        @Override
        public void onPartitionsRevoked(Collection<TopicPartition> parts) {
            currentlyAssigned.removeAll(parts);
        }

        @Override
        public void onPartitionsAssigned(Collection<TopicPartition> parts) {
            currentlyAssigned.addAll(parts);
            log.info("Consumer {} has assigned partitions {}", name, currentlyAssigned);
            emptyPollCount.clear();
            topicPartitionToId.clear();
            AtomicInteger id = new AtomicInteger();
            currentlyAssigned.forEach(p -> {
                topicPartitionToId.put(p, id.getAndIncrement());
                emptyPollCount.put(p, 0);
            });
            if (currentlyAssigned.isEmpty()) {
                watermarkEstimator.set(createWatermarkEstimatorForEmptyParts());
            } else {
                watermarkEstimator.set(new MinimalPartitionWatermarkEstimator(currentlyAssigned.stream().collect(toMap(topicPartitionToId::get, item -> createWatermarkEstimator()))));
            }
            Optional.ofNullable(kafka.get()).ifPresent(c -> consumer.onAssign(c, name != null ? getCommittedTopicOffsets(currentlyAssigned, c) : getCurrentTopicOffsets(currentlyAssigned, c)));
        }

        List<TopicOffset> getCurrentTopicOffsets(Collection<TopicPartition> parts, KafkaConsumer<Object, Object> c) {
            return parts.stream().map(tp -> new TopicOffset(new PartitionWithTopic(tp.topic(), tp.partition()), c.position(tp), watermarkEstimator.get().getWatermark())).collect(Collectors.toList());
        }

        List<TopicOffset> getCommittedTopicOffsets(Collection<TopicPartition> parts, KafkaConsumer<Object, Object> c) {
            Map<TopicPartition, OffsetAndMetadata> committed = new HashMap<>(c.committed(new HashSet<>(parts)));
            for (TopicPartition tp : parts) {
                committed.putIfAbsent(tp, null);
            }
            return committed.entrySet().stream().map(entry -> {
                final long offset = entry.getValue() == null ? 0L : entry.getValue().offset();
                return new TopicOffset(new PartitionWithTopic(entry.getKey().topic(), entry.getKey().partition()), offset, watermarkEstimator.get().getWatermark());
            }).collect(Collectors.toList());
        }

        private WatermarkEstimator createWatermarkEstimator() {
            final WatermarkIdlePolicyFactory idlePolicyFactory = accessor.getWatermarkConfiguration().getWatermarkIdlePolicyFactory();
            final WatermarkEstimatorFactory estimatorFactory = accessor.getWatermarkConfiguration().getWatermarkEstimatorFactory();
            return estimatorFactory.create(cfg, idlePolicyFactory);
        }
    };
}

Example 2 with WatermarkEstimatorFactory

use of cz.o2.proxima.time.WatermarkEstimatorFactory in project proxima-platform by O2-Czech-Republic.

the class PubSubReader method createWatermarkEstimator.

WatermarkEstimator createWatermarkEstimator(long minWatermark) {
    final WatermarkIdlePolicyFactory idlePolicyFactory = watermarkConfiguration.getWatermarkIdlePolicyFactory();
    final WatermarkEstimatorFactory estimatorFactory = watermarkConfiguration.getWatermarkEstimatorFactory();
    final WatermarkEstimator estimator = estimatorFactory.create(cfg, idlePolicyFactory);
    estimator.setMinWatermark(minWatermark);
    return estimator;
}

Example 3 with WatermarkEstimatorFactory

use of cz.o2.proxima.time.WatermarkEstimatorFactory in project proxima-platform by O2-Czech-Republic.

the class BoundedOutOfOrdernessWatermarkEstimatorTest method testFactory.

@Test
public void testFactory() {
    Map<String, Object> cfg = new HashMap<String, Object>() {

        {
            put(prefixedKey(MAX_OUT_OF_ORDERNESS_MS), OUT_OF_ORDERNESS);
        }
    };
    WatermarkIdlePolicyFactory idlePolicyFactory = mock(WatermarkIdlePolicyFactory.class);
    when(idlePolicyFactory.create(cfg)).thenReturn(mock(WatermarkIdlePolicy.class));
    WatermarkEstimatorFactory factory = new BoundedOutOfOrdernessWatermarkEstimator.Factory();
    BoundedOutOfOrdernessWatermarkEstimator watermarkEstimator = (BoundedOutOfOrdernessWatermarkEstimator) factory.create(cfg, idlePolicyFactory);
    assertEquals(OUT_OF_ORDERNESS, watermarkEstimator.getMaxOutOfOrderness());
    verify(idlePolicyFactory, times(1)).create(cfg);
}

Example 4 with WatermarkEstimatorFactory

use of cz.o2.proxima.time.WatermarkEstimatorFactory in project proxima-platform by O2-Czech-Republic.

the class UnboundedOutOfOrdernessWatermarkEstimatorTest method testFactory.

@Test
public void testFactory() {
    long estimateDurationMs = 100L;
    long stepMs = 10L;
    long allowedTimestampSkew = 200L;
    Map<String, Object> cfg = new HashMap<String, Object>() {

        {
            put(prefixedKey(ESTIMATE_DURATION_MS), estimateDurationMs);
            put(prefixedKey(STEP_MS), stepMs);
            put(prefixedKey(ALLOWED_TIMESTAMP_SKEW), allowedTimestampSkew);
        }
    };
    WatermarkIdlePolicyFactory idlePolicyFactory = mock(WatermarkIdlePolicyFactory.class);
    when(idlePolicyFactory.create(cfg)).thenReturn(mock(WatermarkIdlePolicy.class));
    WatermarkEstimatorFactory factory = new UnboundedOutOfOrdernessWatermarkEstimator.Factory();
    UnboundedOutOfOrdernessWatermarkEstimator watermarkEstimator = (UnboundedOutOfOrdernessWatermarkEstimator) factory.create(cfg, idlePolicyFactory);
    assertEquals(estimateDurationMs, watermarkEstimator.getEstimateDurationMs());
    assertEquals(stepMs, watermarkEstimator.getStepMs());
    assertEquals(allowedTimestampSkew, watermarkEstimator.getAllowedTimestampSkew());
    verify(idlePolicyFactory, times(1)).create(cfg);
}