org.apache.kafka.streams.state.KeyValueIterator Java Examples

@Get("/serviceNames")
@ProducesJson
public JsonNode getServiceNames() {
  try {
    if (!storage.traceSearchEnabled) return MAPPER.createArrayNode();
    ReadOnlyKeyValueStore<String, String> store = storage.getTraceStorageStream()
        .store(SERVICE_NAMES_STORE_NAME, QueryableStoreTypes.keyValueStore());
    ArrayNode array = MAPPER.createArrayNode();
    try (KeyValueIterator<String, String> all = store.all()) {
      all.forEachRemaining(keyValue -> array.add(keyValue.value));
    }
    return array;
  } catch (InvalidStateStoreException e) {
    LOG.debug("State store is not ready", e);
    throw e;
  }
}
 

public void cogroup(long timestamp) {
    KeyValueIterator<String, Tuple<List<ClickEvent>, List<StockTransaction>>> iterator = tupleStore.all();

    while (iterator.hasNext()) {
        KeyValue<String, Tuple<List<ClickEvent>, List<StockTransaction>>> cogrouping = iterator.next();

        if (cogrouping.value != null && (!cogrouping.value._1.isEmpty() || !cogrouping.value._2.isEmpty())) {
            List<ClickEvent> clickEvents = new ArrayList<>(cogrouping.value._1);
            List<StockTransaction> stockTransactions = new ArrayList<>(cogrouping.value._2);

            context().forward(cogrouping.key, Tuple.of(clickEvents, stockTransactions));
            cogrouping.value._1.clear();
            cogrouping.value._2.clear();
            tupleStore.put(cogrouping.key, cogrouping.value);
        }
    }
    iterator.close();
}
 

@Override
public void punctuate(long timestamp) {
    KeyValueIterator<String, StockPerformance> performanceIterator = keyValueStore.all();

    while (performanceIterator.hasNext()) {
        KeyValue<String, StockPerformance> keyValue = performanceIterator.next();
        String key = keyValue.key;
        StockPerformance stockPerformance = keyValue.value;

        if (stockPerformance != null) {
            if (stockPerformance.priceDifferential() >= differentialThreshold ||
                    stockPerformance.volumeDifferential() >= differentialThreshold) {
                context.forward(key, stockPerformance);
            }
        }
    }
}
 

@Override
public void punctuate(long timestamp) {
    KeyValueIterator<String, Tuple<List<ClickEvent>, List<StockTransaction>>> iterator = tupleStore.all();

    while (iterator.hasNext()) {
        KeyValue<String, Tuple<List<ClickEvent>, List<StockTransaction>>> cogrouped = iterator.next();
        // if either list contains values forward results
        if (cogrouped.value != null && (!cogrouped.value._1.isEmpty() || !cogrouped.value._2.isEmpty())) {
            List<ClickEvent> clickEvents = new ArrayList<>(cogrouped.value._1);
            List<StockTransaction> stockTransactions = new ArrayList<>(cogrouped.value._2);

            context.forward(cogrouped.key, Tuple.of(clickEvents, stockTransactions));
            // empty out the current cogrouped results
            cogrouped.value._1.clear();
            cogrouped.value._2.clear();
            tupleStore.put(cogrouped.key, cogrouped.value);
        }
    }
    iterator.close();
}
 

@Override
@SuppressWarnings("deprecation")
public KeyValue<String, List<KeyValue<String, StockPerformance>>> punctuate(long timestamp) {
    List<KeyValue<String, StockPerformance>> stockPerformanceList = new ArrayList<>();
    KeyValueIterator<String, StockPerformance> performanceIterator = keyValueStore.all();
    while (performanceIterator.hasNext()) {
        KeyValue<String, StockPerformance> keyValue = performanceIterator.next();
        StockPerformance stockPerformance = keyValue.value;

        if (stockPerformance != null) {
            if (stockPerformance.priceDifferential() >= differentialThreshold ||
                    stockPerformance.volumeDifferential() >= differentialThreshold) {
                stockPerformanceList.add(keyValue);
            }
        }
    }
    return stockPerformanceList.isEmpty() ? null : KeyValue.pair(null, stockPerformanceList);
}
 

private void doPunctuate(long timestamp) {
    KeyValueIterator<String, StockPerformance> performanceIterator = keyValueStore.all();

    while (performanceIterator.hasNext()) {
        KeyValue<String, StockPerformance> keyValue = performanceIterator.next();
        String key = keyValue.key;
        StockPerformance stockPerformance = keyValue.value;

        if (stockPerformance != null) {
            if (stockPerformance.priceDifferential() >= differentialThreshold ||
                    stockPerformance.volumeDifferential() >= differentialThreshold) {
                processorContext.forward(key, stockPerformance);
            }
        }
    }
}
 

@Get("/autocompleteTags")
@ProducesJson
public JsonNode getAutocompleteTags() {
  try {
    if (!storage.traceSearchEnabled) return MAPPER.createArrayNode();
    ReadOnlyKeyValueStore<String, Set<String>> autocompleteTagsStore =
        storage.getTraceStorageStream().store(AUTOCOMPLETE_TAGS_STORE_NAME,
            QueryableStoreTypes.keyValueStore());
    ArrayNode array = MAPPER.createArrayNode();
    try (KeyValueIterator<String, Set<String>> all = autocompleteTagsStore.all()) {
      all.forEachRemaining(keyValue -> array.add(keyValue.key));
    }
    return array;
  } catch (InvalidStateStoreException e) {
    LOG.debug("State store is not ready", e);
    throw e;
  }
}
 

void addResults(
    QueryRequest request,
    ReadOnlyKeyValueStore<String, List<Span>> tracesStore,
    List<List<Span>> traces,
    List<String> traceIds,
    KeyValueIterator<Long, Set<String>> spanIds
) {
  spanIds.forEachRemaining(keyValue -> {
    for (String traceId : keyValue.value) {
      if (!traceIds.contains(traceId)) {
        List<Span> spans = tracesStore.get(traceId);
        if (spans != null && !spans.isEmpty() && request.test(spans)) { // apply filters
          traceIds.add(traceId); // adding to check if we have already add it later
          traces.add(spans);
        }
      }
    }
  });
}
 

/**
 * Query local state store to extract metrics
 *
 * @return local Metrics
 */
private Metrics getLocalMetrics() {
    HostInfo thisInstance = GlobalAppState.getInstance().getHostPortInfo();
    KafkaStreams ks = GlobalAppState.getInstance().getKafkaStreams();

    String source = thisInstance.host() + ":" + thisInstance.port();
    Metrics localMetrics = new Metrics();

    ReadOnlyKeyValueStore<String, Double> averageStore = ks
            .store(storeName,
                    QueryableStoreTypes.<String, Double>keyValueStore());

    LOGGER.log(Level.INFO, "Entries in store {0}", averageStore.approximateNumEntries());
    KeyValueIterator<String, Double> storeIterator = averageStore.all();

    while (storeIterator.hasNext()) {
        KeyValue<String, Double> kv = storeIterator.next();
        localMetrics.add(source, kv.key, String.valueOf(kv.value));

    }
    LOGGER.log(Level.INFO, "Local store state {0}", localMetrics);
    return localMetrics;
}
 

@Override
public void start() throws Exception {

    execute(Config.RANGE_KEY_VALUE_QUERY_ADDRESS_PREFIX, (abstractQuery, keySerde, valueSerde) -> {

        RangeKeyValueQuery query = (RangeKeyValueQuery) abstractQuery;
        ReadOnlyKeyValueStore<Object, Object> kvStore = streams.store(query.getStoreName(), QueryableStoreTypes.keyValueStore());
        MultiValuedKeyValueQueryResponse response;
        try (KeyValueIterator<Object, Object> result = kvStore.range(deserializeObject(keySerde, query.getFrom()), deserializeObject(keySerde, query.getTo()))) {
            if (result.hasNext()) {
                Map<String, String> results = new HashMap<>();
                while (result.hasNext()) {
                    KeyValue<Object, Object> kvEntry = result.next();
                    results.put(base64Encode(keySerde, kvEntry.key), base64Encode(valueSerde, kvEntry.value));
                }
                return new MultiValuedKeyValueQueryResponse(results);
            } else {
               return new MultiValuedKeyValueQueryResponse(Collections.emptyMap());
            }
        }
    } );


}
 

@Override
public void start() throws Exception {

    execute(Config.ALL_KEY_VALUE_QUERY_ADDRESS_PREFIX, (query, keySerde, valueSerde) -> {

        ReadOnlyKeyValueStore<Object, Object> kvStore = streams.store(query.getStoreName(), QueryableStoreTypes.keyValueStore());
        MultiValuedKeyValueQueryResponse response;
        try (KeyValueIterator<Object, Object> result = kvStore.all()) {
            if (result.hasNext()) {
                Map<String, String> results = new HashMap<>();
                while (result.hasNext()) {
                    KeyValue<Object, Object> kvEntry = result.next();

                    results.put(base64Encode(keySerde, kvEntry.key), base64Encode(valueSerde, kvEntry.value));
                }
               return new MultiValuedKeyValueQueryResponse(results);
            } else {
               return new MultiValuedKeyValueQueryResponse(Collections.emptyMap());
            }
        }
    });

}
 

@Test
public void notFoundWithNoResult(TestContext context){
    KafkaStreams streamMock = mock(KafkaStreams.class);
    ReadOnlyKeyValueStore<Object, Object> storeMock = mock(ReadOnlyKeyValueStore.class);
    KeyValueIterator<Object, Object> iteratorMock = mock(KeyValueIterator.class);
    when(streamMock.store(eq("store"), any(QueryableStoreType.class))).thenReturn(storeMock);
    SimpleKeyValueIterator iterator = new SimpleKeyValueIterator();
    when(storeMock.all()).thenReturn(iterator);


    rule.vertx().deployVerticle(new AllKeyValuesQueryVerticle("host", streamMock), context.asyncAssertSuccess(deployment->{

        StoreWideQuery query = new StoreWideQuery("store", Serdes.String().getClass().getName(), Serdes.String().getClass().getName());

        rule.vertx().eventBus().send(Config.ALL_KEY_VALUE_QUERY_ADDRESS_PREFIX + "host", query, context.asyncAssertSuccess(reply ->{

            context.assertTrue(reply.body() instanceof MultiValuedKeyValueQueryResponse);
            MultiValuedKeyValueQueryResponse response = (MultiValuedKeyValueQueryResponse) reply.body();
            context.assertEquals(0, response.getResults().size());
            context.assertTrue(iterator.closed);

        }));

    }));

}
 

@Test
public void notFoundWithNoResult(TestContext context){
    KafkaStreams streamMock = mock(KafkaStreams.class);
    ReadOnlyKeyValueStore<Object, Object> storeMock = mock(ReadOnlyKeyValueStore.class);
    KeyValueIterator<Object, Object> iteratorMock = mock(KeyValueIterator.class);
    when(streamMock.store(eq("store"), any(QueryableStoreType.class))).thenReturn(storeMock);
    SimpleKeyValueIterator iterator = new SimpleKeyValueIterator();
    when(storeMock.range(any(), any())).thenReturn(iterator);


    rule.vertx().deployVerticle(new RangeKeyValueQueryVerticle("host", streamMock), context.asyncAssertSuccess(deployment->{

        RangeKeyValueQuery query = new RangeKeyValueQuery("store", Serdes.String().getClass().getName(), Serdes.String().getClass().getName(), "key".getBytes(), "key".getBytes());

        rule.vertx().eventBus().send(Config.RANGE_KEY_VALUE_QUERY_ADDRESS_PREFIX + "host", query, context.asyncAssertSuccess(reply ->{

            context.assertTrue(reply.body() instanceof MultiValuedKeyValueQueryResponse);
            MultiValuedKeyValueQueryResponse response = (MultiValuedKeyValueQueryResponse) reply.body();
            context.assertEquals(0, response.getResults().size());
            context.assertTrue(iterator.closed);

        }));

    }));

}
 

@Override
public KeyValueIterator<K, V> range(K from, K to) {
    ByteString fromBytes = ByteString.copyFrom(keyValueSerde.serializeKey(from));
    ByteString toBytes = ByteString.copyFrom(keyValueSerde.serializeKey(to));
    StreamObserverSpliterator<io.apicurio.registry.streams.distore.proto.KeyValue> observer = new StreamObserverSpliterator<>();
    stub.range(
        KeyFromKeyToReq
            .newBuilder()
            .setKeyFrom(fromBytes)
            .setKeyTo(toBytes)
            .setStoreName(storeName)
            .build(),
        observer
    );
    return keyValueIterator(observer.stream());
}
 

@Override
public void all(VoidReq request, StreamObserver<io.apicurio.registry.streams.distore.proto.KeyValue> responseObserver) {
    boolean ok = false;
    try (
        KeyValueIterator<?, ?> iter =
            keyValueStore(request.getStoreName()).all()
    ) {
        drainToKeyValue(request.getStoreName(), iter, responseObserver);
        ok = true;
    } catch (Throwable e) {
        responseObserver.onError(e);
    }
    if (ok) {
        responseObserver.onCompleted();
    }
}
 

@Override
public void range(KeyFromKeyToReq request, StreamObserver<io.apicurio.registry.streams.distore.proto.KeyValue> responseObserver) {
    boolean ok = false;
    try (
        KeyValueIterator<?, ?> iter =
            keyValueStore(request.getStoreName()).range(
                keyValueSerdes.deserializeKey(request.getStoreName(), request.getKeyFrom().toByteArray()),
                keyValueSerdes.deserializeVal(request.getStoreName(), request.getKeyTo().toByteArray())
            )
    ) {
        drainToKeyValue(request.getStoreName(), iter, responseObserver);
        ok = true;
    } catch (Throwable e) {
        responseObserver.onError(e);
    }
    if (ok) {
        responseObserver.onCompleted();
    }
}
 

private static void printStateStoreKeyValueIterator(final KeyValueIterator<String, VisibilityBindingSet> rangeIt) {
    log.info("----------------");
    while (rangeIt.hasNext()) {
        final KeyValue<String, VisibilityBindingSet> keyValue = rangeIt.next();
        log.info(keyValue.key + " :::: " + keyValue.value);
    }
    log.info("----------------\n\n");
    if (rangeIt != null) {
        rangeIt.close();
    }
}
 

@Override
public void punctuate(long timestamp) {
    CuratorFramework curator = CuratorContainer.getInstance().getCurator();
    KeyValueIterator<String, String> iter = store.all();
    while (iter.hasNext()) {
        KeyValue<String, String> entry = iter.next();
        try {
            curator.setData().forPath(entry.key, entry.value.getBytes());
        } catch (Exception e) {
        }
    }
}
 

public static <K, V> NavigableMap<K, V> mapFromStore(KafkaStreams streams, String storeName) {
    final ReadOnlyKeyValueStore<K, V> store = streams.store(
        storeName, QueryableStoreTypes.keyValueStore());

    try (final KeyValueIterator<K, V> all = store.all()) {
        NavigableMap<K, V> result = new TreeMap<>();
        while (all.hasNext()) {
            KeyValue<K, V> next = all.next();
            result.put(next.key, next.value);
        }
        return result;
    }
}
 

public static <K, V> List<KeyValue<K, V>> listFromStore(KafkaStreams streams, String storeName) {
    final ReadOnlyKeyValueStore<K, V> store = streams.store(
        storeName, QueryableStoreTypes.keyValueStore());

    try (final KeyValueIterator<K, V> all = store.all()) {
        List<KeyValue<K, V>> result = new ArrayList<>();
        while (all.hasNext()) {
            result.add(all.next());
        }
        return result;
    }
}
 

@Override
public KeyValueIterator<K, V> all() {
    List<KeyValue<K, V>> entryList = new ArrayList<>();
    for (Map.Entry<K, V> tupleEntry : inner.entrySet()) {
        entryList.add(KeyValue.pair(tupleEntry.getKey(),tupleEntry.getValue()));
    }
    return new KeyValueIteratorStub<>(entryList.iterator());
}
 

public List<Customer> getCustomers() {
    List<Customer> customers = new ArrayList<>();
    KeyValueIterator<UUID, Map> iterator = store.all();
    while (iterator.hasNext()) {
        KeyValue<UUID, Map> entry = iterator.next();
        logger.debug("getCustomers iterator entry: {}", entry);
        customers.add(new Customer(entry.value));
    }
    iterator.close();
    return customers;
}
 

@Override
public void punctuate(long streamTime) {
    KeyValueIterator<String, StockTransactionSummary> it = summaryStore.all();
    long currentTime = System.currentTimeMillis();
    while (it.hasNext()) {
        StockTransactionSummary summary = it.next().value;
        if (summary.updatedWithinLastMillis(currentTime, 11000)) {
            this.context.forward(summary.tickerSymbol, summary);
        }
    }
}
 

@Override
public void punctuate(long timestamp) {
	try (KeyValueIterator<String, Integer> iterator = this.kvStore.all()) {
		iterator.forEachRemaining(entry -> {
			context.forward(entry.key, entry.value);
			this.kvStore.delete(entry.key);
		});
	}
	context.commit();
}
 

private String getKeyValuesAsJson(String store) {
    ReadOnlyKeyValueStore<String, Long> readOnlyKeyValueStore = kafkaStreams.store(store, QueryableStoreTypes.keyValueStore());
    List<KeyValue<String, Long>> keyValues = new ArrayList<>();
    try (KeyValueIterator<String, Long> iterator = readOnlyKeyValueStore.all()) {
        while (iterator.hasNext()) {
            keyValues.add(iterator.next());
        }
    }
    return gson.toJson(keyValues);
}
 

private String fetchFromSessionStore(Map<String, String> params) {
    String store = params.get(STORE_PARAM);
    String key = params.get(KEY_PARAM);

    HostInfo storeHostInfo = getHostInfo(store, key);

    if (storeHostInfo.host().equals("unknown")) {
        return STORES_NOT_ACCESSIBLE;
    }

    if (dataNotLocal(storeHostInfo)) {
        LOG.info("{} located in state store on another instance !!!!", key);
        return fetchRemote(storeHostInfo, "session", params);
    }

    ReadOnlySessionStore<String, CustomerTransactions> readOnlySessionStore = kafkaStreams.store(store, QueryableStoreTypes.sessionStore());
    List<String> results = new ArrayList<>();
    List<KeyValue<String, List<String>>> sessionResults = new ArrayList<>();
    try (KeyValueIterator<Windowed<String>, CustomerTransactions> iterator = readOnlySessionStore.fetch(key)) {
        while (iterator.hasNext()) {
            KeyValue<Windowed<String>, CustomerTransactions> windowed = iterator.next();
            CustomerTransactions transactions = windowed.value;
            LocalDateTime startSession = getLocalDateTime(windowed.key.window().start());
            LocalDateTime endSession = getLocalDateTime(windowed.key.window().end());
            transactions.setSessionInfo(String.format("Session Window{start=%s, end=%s}", startSession.toLocalTime().toString(), endSession.toLocalTime().toString()));
            results.add(transactions.toString());
        }
        sessionResults.add(new KeyValue<>(key, results));
    }

    return gson.toJson(sessionResults);
}
 

@Get("/dependencies")
public AggregatedHttpResponse getDependencies(
    @Param("endTs") long endTs,
    @Param("lookback") long lookback
) {
  try {
    if (!storage.dependencyQueryEnabled) return AggregatedHttpResponse.of(HttpStatus.NOT_FOUND);
    ReadOnlyWindowStore<Long, DependencyLink> store =
        storage.getDependencyStorageStream()
            .store(DEPENDENCIES_STORE_NAME, QueryableStoreTypes.windowStore());
    List<DependencyLink> links = new ArrayList<>();
    Instant from = Instant.ofEpochMilli(endTs - lookback);
    Instant to = Instant.ofEpochMilli(endTs);
    try (KeyValueIterator<Windowed<Long>, DependencyLink> iterator = store.fetchAll(from, to)) {
      iterator.forEachRemaining(keyValue -> links.add(keyValue.value));
    }
    List<DependencyLink> mergedLinks = DependencyLinker.merge(links);
    LOG.debug("Dependencies found from={}-to={}: {}", from, to, mergedLinks.size());
    return AggregatedHttpResponse.of(
        HttpStatus.OK,
        MediaType.JSON,
        DependencyLinkBytesEncoder.JSON_V1.encodeList(mergedLinks));
  } catch (InvalidStateStoreException e) {
    LOG.debug("State store is not ready", e);
    return AggregatedHttpResponse.of(HttpStatus.SERVICE_UNAVAILABLE);
  }
}
 

@Override
public KeyValueIterator<K, V> all() {
    return new StreamToKeyValueIteratorAdapter<>(
        allServicesForStoreStream()
            .flatMap(store -> toStream(store.all()))
    );
}
 

@Override
public KeyValueIterator<K, V> range(K from, K to) {
    return new StreamToKeyValueIteratorAdapter<>(
        allServicesForStoreStream()
            .flatMap(store -> toStream(store.range(from, to)))
    );
}
 

public static <K, V> Stream<KeyValue<K, V>> toStream(KeyValueIterator<K, V> kvIterator) {
    if (kvIterator instanceof StreamToKeyValueIteratorAdapter) {
        return ((StreamToKeyValueIteratorAdapter<K, V>) kvIterator).stream;
    }
    return StreamSupport
        .stream(
            Spliterators.spliteratorUnknownSize(
                kvIterator,
                Spliterator.ORDERED | Spliterator.IMMUTABLE | Spliterator.NONNULL
            ),
            false
        )
        .onClose(kvIterator::close);
}