diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/CachingSessionStore.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/CachingSessionStore.java index 022f6f36ddd35..cb0cb250a7e94 100644 --- a/streams/src/main/java/org/apache/kafka/streams/state/internals/CachingSessionStore.java +++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/CachingSessionStore.java @@ -143,7 +143,7 @@ public void put(final Windowed key, byte[] value) { validateStoreOpen(); final Bytes binaryKey = SessionKeySerde.bytesToBinary(key); final LRUCacheEntry entry = new LRUCacheEntry(value, true, context.offset(), - key.window().end(), context.partition(), context.topic()); + context.timestamp(), context.partition(), context.topic()); cache.put(cacheName, cacheFunction.cacheKey(binaryKey), entry); } diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/CachingWindowStore.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/CachingWindowStore.java index a78978b41c1bc..99c3e7f141610 100644 --- a/streams/src/main/java/org/apache/kafka/streams/state/internals/CachingWindowStore.java +++ b/streams/src/main/java/org/apache/kafka/streams/state/internals/CachingWindowStore.java @@ -152,7 +152,7 @@ public synchronized void put(final Bytes key, final byte[] value, final long tim final Bytes keyBytes = WindowStoreUtils.toBinaryKey(key, timestamp, 0, bytesSerdes); final LRUCacheEntry entry = new LRUCacheEntry(value, true, context.offset(), - timestamp, context.partition(), context.topic()); + context.timestamp(), context.partition(), context.topic()); cache.put(name, cacheFunction.cacheKey(keyBytes), entry); } diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/KStreamAggregationIntegrationTest.java b/streams/src/test/java/org/apache/kafka/streams/integration/KStreamAggregationIntegrationTest.java index 4527c19b471ed..3afded51d46f8 100644 --- a/streams/src/test/java/org/apache/kafka/streams/integration/KStreamAggregationIntegrationTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/integration/KStreamAggregationIntegrationTest.java @@ -45,11 +45,16 @@ import org.apache.kafka.streams.kstream.Serialized; import org.apache.kafka.streams.kstream.SessionWindows; import org.apache.kafka.streams.kstream.TimeWindows; +import org.apache.kafka.streams.kstream.Transformer; +import org.apache.kafka.streams.kstream.TransformerSupplier; import org.apache.kafka.streams.kstream.Windowed; import org.apache.kafka.streams.kstream.internals.SessionWindow; +import org.apache.kafka.streams.processor.ProcessorContext; import org.apache.kafka.streams.state.KeyValueIterator; +import org.apache.kafka.streams.state.KeyValueStore; import org.apache.kafka.streams.state.QueryableStoreTypes; import org.apache.kafka.streams.state.ReadOnlySessionStore; +import org.apache.kafka.streams.state.SessionStore; import org.apache.kafka.streams.state.WindowStore; import org.apache.kafka.test.IntegrationTest; import org.apache.kafka.test.MockMapper; @@ -196,6 +201,15 @@ private static int compare(final Ke return keyComparison; } + private static int compareIgnoreTimestamp(final KeyValue> o1, + final KeyValue> o2) { + final int keyComparison = o1.key.compareTo(o2.key); + if (keyComparison == 0) { + return o1.value.key.compareTo(o2.value.key); + } + return keyComparison; + } + @Test public void shouldReduceWindowed() throws Exception { final long firstBatchTimestamp = mockTime.milliseconds(); @@ -325,18 +339,18 @@ public String apply(final Windowed windowedKey, final Integer value) { startStreams(); - final List> windowedMessages = receiveMessages( + final List>> windowedMessages = receiveMessagesWithTimestamp( new StringDeserializer(), new IntegerDeserializer(), 15); - final Comparator> + final Comparator>> comparator = - new Comparator>() { + new Comparator>>() { @Override - public int compare(final KeyValue o1, - final KeyValue o2) { - return KStreamAggregationIntegrationTest.compare(o1, o2); + public int compare(final KeyValue> o1, + final KeyValue> o2) { + return KStreamAggregationIntegrationTest.compareIgnoreTimestamp(o1, o2); } }; @@ -347,21 +361,21 @@ public int compare(final KeyValue o1, assertThat(windowedMessages, is( Arrays.asList( - new KeyValue<>("A@" + firstWindow, 1), - new KeyValue<>("A@" + secondWindow, 1), - new KeyValue<>("A@" + secondWindow, 2), - new KeyValue<>("B@" + firstWindow, 1), - new KeyValue<>("B@" + secondWindow, 1), - new KeyValue<>("B@" + secondWindow, 2), - new KeyValue<>("C@" + firstWindow, 1), - new KeyValue<>("C@" + secondWindow, 1), - new KeyValue<>("C@" + secondWindow, 2), - new KeyValue<>("D@" + firstWindow, 1), - new KeyValue<>("D@" + secondWindow, 1), - new KeyValue<>("D@" + secondWindow, 2), - new KeyValue<>("E@" + firstWindow, 1), - new KeyValue<>("E@" + secondWindow, 1), - new KeyValue<>("E@" + secondWindow, 2) + new KeyValue<>("A@" + firstWindow, KeyValue.pair(1, firstTimestamp)), + new KeyValue<>("A@" + secondWindow, KeyValue.pair(1, secondTimestamp)), + new KeyValue<>("A@" + secondWindow, KeyValue.pair(2, secondTimestamp)), + new KeyValue<>("B@" + firstWindow, KeyValue.pair(1, firstTimestamp)), + new KeyValue<>("B@" + secondWindow, KeyValue.pair(1, secondTimestamp)), + new KeyValue<>("B@" + secondWindow, KeyValue.pair(2, secondTimestamp)), + new KeyValue<>("C@" + firstWindow, KeyValue.pair(1, firstTimestamp)), + new KeyValue<>("C@" + secondWindow, KeyValue.pair(1, secondTimestamp)), + new KeyValue<>("C@" + secondWindow, KeyValue.pair(2, secondTimestamp)), + new KeyValue<>("D@" + firstWindow, KeyValue.pair(1, firstTimestamp)), + new KeyValue<>("D@" + secondWindow, KeyValue.pair(1, secondTimestamp)), + new KeyValue<>("D@" + secondWindow, KeyValue.pair(2, secondTimestamp)), + new KeyValue<>("E@" + firstWindow, KeyValue.pair(1, firstTimestamp)), + new KeyValue<>("E@" + secondWindow, KeyValue.pair(1, secondTimestamp)), + new KeyValue<>("E@" + secondWindow, KeyValue.pair(2, secondTimestamp)) ))); } @@ -400,8 +414,9 @@ public int compare(final KeyValue o1, final KeyValue public void shouldCount() throws Exception { produceMessages(mockTime.milliseconds()); - groupedStream.count("count-by-key") - .to(Serdes.String(), Serdes.Long(), outputTopic); + groupedStream.count(Materialized.>as("count-by-key")) + .toStream() + .to(outputTopic, Produced.with(Serdes.String(), Serdes.Long())); shouldCountHelper(); } @@ -524,30 +539,51 @@ public void shouldCountSessionWindows() throws Exception { new Properties()), t4); - final Map, Long> results = new HashMap<>(); + final Map, KeyValue> results = new HashMap<>(); final CountDownLatch latch = new CountDownLatch(11); builder.stream(userSessionsStream, Consumed.with(Serdes.String(), Serdes.String())) .groupByKey(Serialized.with(Serdes.String(), Serdes.String())) .count(SessionWindows.with(sessionGap).until(maintainMillis)) .toStream() - .foreach(new ForeachAction, Long>() { + .transform(new TransformerSupplier, Long, KeyValue>() { @Override - public void apply(final Windowed key, final Long value) { - results.put(key, value); - latch.countDown(); + public Transformer, Long, KeyValue> get() { + return new Transformer, Long, KeyValue>() { + private ProcessorContext context; + + @Override + public void init(final ProcessorContext context) { + this.context = context; + } + + @Override + public KeyValue transform(final Windowed key, final Long value) { + results.put(key, KeyValue.pair(value, context.timestamp())); + latch.countDown(); + return null; + } + + @Override + public KeyValue punctuate(final long timestamp) { + return null; + } + + @Override + public void close() {} + }; } }); startStreams(); latch.await(30, TimeUnit.SECONDS); - assertThat(results.get(new Windowed<>("bob", new SessionWindow(t1, t1))), equalTo(1L)); - assertThat(results.get(new Windowed<>("penny", new SessionWindow(t1, t1))), equalTo(1L)); - assertThat(results.get(new Windowed<>("jo", new SessionWindow(t1, t1))), equalTo(1L)); - assertThat(results.get(new Windowed<>("jo", new SessionWindow(t4, t4))), equalTo(1L)); - assertThat(results.get(new Windowed<>("emily", new SessionWindow(t1, t2))), equalTo(2L)); - assertThat(results.get(new Windowed<>("bob", new SessionWindow(t3, t4))), equalTo(2L)); - assertThat(results.get(new Windowed<>("penny", new SessionWindow(t3, t3))), equalTo(1L)); + assertThat(results.get(new Windowed<>("bob", new SessionWindow(t1, t1))), equalTo(KeyValue.pair(1L, t1))); + assertThat(results.get(new Windowed<>("penny", new SessionWindow(t1, t1))), equalTo(KeyValue.pair(1L, t1))); + assertThat(results.get(new Windowed<>("jo", new SessionWindow(t1, t1))), equalTo(KeyValue.pair(1L, t1))); + assertThat(results.get(new Windowed<>("jo", new SessionWindow(t4, t4))), equalTo(KeyValue.pair(1L, t4))); + assertThat(results.get(new Windowed<>("emily", new SessionWindow(t1, t2))), equalTo(KeyValue.pair(2L, t2))); + assertThat(results.get(new Windowed<>("bob", new SessionWindow(t3, t4))), equalTo(KeyValue.pair(2L, t4))); + assertThat(results.get(new Windowed<>("penny", new SessionWindow(t3, t3))), equalTo(KeyValue.pair(1L, t3))); } @SuppressWarnings("deprecation") @@ -617,19 +653,20 @@ public void shouldReduceSessionWindows() throws Exception { final String userSessionsStore = "UserSessionsStore"; builder.stream(userSessionsStream, Consumed.with(Serdes.String(), Serdes.String())) .groupByKey(Serialized.with(Serdes.String(), Serdes.String())) + .windowedBy(SessionWindows.with(sessionGap).until(maintainMillis)) .reduce(new Reducer() { @Override public String apply(final String value1, final String value2) { return value1 + ":" + value2; } - }, SessionWindows.with(sessionGap).until(maintainMillis), userSessionsStore) - .foreach(new ForeachAction, String>() { - @Override - public void apply(final Windowed key, final String value) { - results.put(key, value); - latch.countDown(); - } - }); + }, Materialized.>as(userSessionsStore).withValueSerde(Serdes.String())) + .foreach(new ForeachAction, String>() { + @Override + public void apply(final Windowed key, final String value) { + results.put(key, value); + latch.countDown(); + } + }); startStreams(); latch.await(30, TimeUnit.SECONDS); @@ -650,10 +687,8 @@ public void apply(final Windowed key, final String value) { assertThat(bob.next(), equalTo(KeyValue.pair(new Windowed<>("bob", new SessionWindow(t1, t1)), "start"))); assertThat(bob.next(), equalTo(KeyValue.pair(new Windowed<>("bob", new SessionWindow(t3, t4)), "pause:resume"))); assertFalse(bob.hasNext()); - } - private void produceMessages(final long timestamp) throws Exception { IntegrationTestUtils.produceKeyValuesSynchronouslyWithTimestamp( streamOneInput, @@ -671,7 +706,6 @@ private void produceMessages(final long timestamp) throws Exception { timestamp); } - private void createTopics() throws InterruptedException { streamOneInput = "stream-one-" + testNo; outputTopic = "output-" + testNo; @@ -685,16 +719,13 @@ private void startStreams() { kafkaStreams.start(); } - - private List> receiveMessages(final Deserializer - keyDeserializer, - final Deserializer - valueDeserializer, + private List> receiveMessages(final Deserializer keyDeserializer, + final Deserializer valueDeserializer, final int numMessages) throws InterruptedException { + final Properties consumerProperties = new Properties(); - consumerProperties - .setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, CLUSTER.bootstrapServers()); + consumerProperties.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, CLUSTER.bootstrapServers()); consumerProperties.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "kgroupedstream-test-" + testNo); consumerProperties.setProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest"); consumerProperties.setProperty(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, keyDeserializer.getClass().getName()); @@ -704,7 +735,24 @@ private List> receiveMessages(final Deserializer outputTopic, numMessages, 60 * 1000); + } + private List>> receiveMessagesWithTimestamp(final Deserializer keyDeserializer, + final Deserializer valueDeserializer, + final int numMessages) + throws InterruptedException { + + final Properties consumerProperties = new Properties(); + consumerProperties.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, CLUSTER.bootstrapServers()); + consumerProperties.setProperty(ConsumerConfig.GROUP_ID_CONFIG, "kgroupedstream-test-" + testNo); + consumerProperties.setProperty(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest"); + consumerProperties.setProperty(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, keyDeserializer.getClass().getName()); + consumerProperties.setProperty(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, valueDeserializer.getClass().getName()); + return IntegrationTestUtils.waitUntilMinKeyValueWithTimestampRecordsReceived( + consumerProperties, + outputTopic, + numMessages, + 60 * 1000); } } diff --git a/streams/src/test/java/org/apache/kafka/streams/integration/utils/IntegrationTestUtils.java b/streams/src/test/java/org/apache/kafka/streams/integration/utils/IntegrationTestUtils.java index 304a3e55db7d4..10d408316110a 100644 --- a/streams/src/test/java/org/apache/kafka/streams/integration/utils/IntegrationTestUtils.java +++ b/streams/src/test/java/org/apache/kafka/streams/integration/utils/IntegrationTestUtils.java @@ -217,6 +217,37 @@ public boolean conditionMet() { return accumData; } + /** + * Wait until enough data (key-value records) has been consumed. + * + * @param consumerConfig Kafka Consumer configuration + * @param topic Topic to consume from + * @param expectedNumRecords Minimum number of expected records + * @param waitTime Upper bound in waiting time in milliseconds + * @return All the records consumed, or null if no records are consumed + * @throws AssertionError if the given wait time elapses + */ + public static List>> waitUntilMinKeyValueWithTimestampRecordsReceived(final Properties consumerConfig, + final String topic, + final int expectedNumRecords, + final long waitTime) throws InterruptedException { + final List>> accumData = new ArrayList<>(); + try (final Consumer consumer = createConsumer(consumerConfig)) { + final TestCondition valuesRead = new TestCondition() { + @Override + public boolean conditionMet() { + final List>> readData = + readKeyValuesWithTimestamp(topic, consumer, waitTime, expectedNumRecords); + accumData.addAll(readData); + return accumData.size() >= expectedNumRecords; + } + }; + final String conditionDetails = "Did not receive all " + expectedNumRecords + " records from topic " + topic; + TestUtils.waitForCondition(valuesRead, waitTime, conditionDetails); + } + return accumData; + } + public static List waitUntilMinValuesRecordsReceived(final Properties consumerConfig, final String topic, final int expectedNumRecords) throws InterruptedException { @@ -382,6 +413,37 @@ private static List> readKeyValues(final String topic, return consumedValues; } + /** + * Returns up to `maxMessages` by reading via the provided consumer (the topic(s) to read from + * are already configured in the consumer). + * + * @param topic Kafka topic to read messages from + * @param consumer Kafka consumer + * @param waitTime Maximum wait time in milliseconds + * @param maxMessages Maximum number of messages to read via the consumer + * @return The KeyValue elements retrieved via the consumer + */ + private static List>> readKeyValuesWithTimestamp(final String topic, + final Consumer consumer, + final long waitTime, + final int maxMessages) { + final List>> consumedValues; + consumer.subscribe(Collections.singletonList(topic)); + final int pollIntervalMs = 100; + consumedValues = new ArrayList<>(); + int totalPollTimeMs = 0; + while (totalPollTimeMs < waitTime && + continueConsuming(consumedValues.size(), maxMessages)) { + totalPollTimeMs += pollIntervalMs; + final ConsumerRecords records = consumer.poll(pollIntervalMs); + + for (final ConsumerRecord record : records) { + consumedValues.add(new KeyValue<>(record.key(), KeyValue.pair(record.value(), record.timestamp()))); + } + } + return consumedValues; + } + private static boolean continueConsuming(final int messagesConsumed, final int maxMessages) { return maxMessages <= 0 || messagesConsumed < maxMessages; }