Examples with StateRestoreListener org.apache.kafka.streams.processor.StateRestoreListener used on opensource projects

Example 1 with StateRestoreListener

use of org.apache.kafka.streams.processor.StateRestoreListener in project apache-kafka-on-k8s by banzaicloud.

the class InternalMockProcessorContext method restore.

public void restore(final String storeName, final Iterable<KeyValue<byte[], byte[]>> changeLog) {
    final BatchingStateRestoreCallback restoreCallback = getBatchingRestoreCallback(restoreFuncs.get(storeName));
    final StateRestoreListener restoreListener = getStateRestoreListener(restoreCallback);
    restoreListener.onRestoreStart(null, storeName, 0L, 0L);
    List<KeyValue<byte[], byte[]>> records = new ArrayList<>();
    for (KeyValue<byte[], byte[]> keyValue : changeLog) {
        records.add(keyValue);
    }
    restoreCallback.restoreAll(records);
    restoreListener.onRestoreEnd(null, storeName, 0L);
}

Example 2 with StateRestoreListener

use of org.apache.kafka.streams.processor.StateRestoreListener in project apache-kafka-on-k8s by banzaicloud.

the class RestoreIntegrationTest method shouldRestoreState.

@Test
public void shouldRestoreState() throws ExecutionException, InterruptedException {
    final AtomicInteger numReceived = new AtomicInteger(0);
    final StreamsBuilder builder = new StreamsBuilder();
    createStateForRestoration();
    builder.table(INPUT_STREAM, Consumed.with(Serdes.Integer(), Serdes.Integer())).toStream().foreach(new ForeachAction<Integer, Integer>() {

        @Override
        public void apply(final Integer key, final Integer value) {
            numReceived.incrementAndGet();
        }
    });
    final CountDownLatch startupLatch = new CountDownLatch(1);
    kafkaStreams = new KafkaStreams(builder.build(), props(applicationId));
    kafkaStreams.setStateListener(new KafkaStreams.StateListener() {

        @Override
        public void onChange(final KafkaStreams.State newState, final KafkaStreams.State oldState) {
            if (newState == KafkaStreams.State.RUNNING && oldState == KafkaStreams.State.REBALANCING) {
                startupLatch.countDown();
            }
        }
    });
    final AtomicLong restored = new AtomicLong(0);
    kafkaStreams.setGlobalStateRestoreListener(new StateRestoreListener() {

        @Override
        public void onRestoreStart(final TopicPartition topicPartition, final String storeName, final long startingOffset, final long endingOffset) {
        }

        @Override
        public void onBatchRestored(final TopicPartition topicPartition, final String storeName, final long batchEndOffset, final long numRestored) {
        }

        @Override
        public void onRestoreEnd(final TopicPartition topicPartition, final String storeName, final long totalRestored) {
            restored.addAndGet(totalRestored);
        }
    });
    kafkaStreams.start();
    assertTrue(startupLatch.await(30, TimeUnit.SECONDS));
    assertThat(restored.get(), equalTo((long) numberOfKeys));
    assertThat(numReceived.get(), equalTo(0));
}