From 4d0e23f78780db4d337797326237807e760ba429 Mon Sep 17 00:00:00 2001 From: Bill Bejeck Date: Wed, 24 Jul 2019 15:57:42 -0400 Subject: [PATCH] KAFKA-8705: Remove parent node after leaving loop to prevent NPE --- .../internals/InternalStreamsBuilder.java | 9 +++- .../internals/graph/StreamsGraphTest.java | 49 +++++++++++++++++++ 2 files changed, 56 insertions(+), 2 deletions(-) diff --git a/streams/src/main/java/org/apache/kafka/streams/kstream/internals/InternalStreamsBuilder.java b/streams/src/main/java/org/apache/kafka/streams/kstream/internals/InternalStreamsBuilder.java index 3a90fd222d484..e48a3fc538050 100644 --- a/streams/src/main/java/org/apache/kafka/streams/kstream/internals/InternalStreamsBuilder.java +++ b/streams/src/main/java/org/apache/kafka/streams/kstream/internals/InternalStreamsBuilder.java @@ -42,6 +42,7 @@ import java.util.Collections; import java.util.Comparator; import java.util.HashMap; +import java.util.HashSet; import java.util.Iterator; import java.util.LinkedHashMap; import java.util.LinkedHashSet; @@ -378,6 +379,7 @@ private void maybeOptimizeRepartitionOperations() { private void maybeUpdateKeyChangingRepartitionNodeMap() { final Map> mergeNodesToKeyChangers = new HashMap<>(); + final Set mergeNodeKeyChangingParentsToRemove = new HashSet<>(); for (final StreamsGraphNode mergeNode : mergeNodes) { mergeNodesToKeyChangers.put(mergeNode, new LinkedHashSet<>()); final Collection keys = keyChangingOperationsToOptimizableRepartitionNodes.keySet(); @@ -395,11 +397,14 @@ private void maybeUpdateKeyChangingRepartitionNodeMap() { final LinkedHashSet repartitionNodes = new LinkedHashSet<>(); for (final StreamsGraphNode keyChangingParent : keyChangingParents) { repartitionNodes.addAll(keyChangingOperationsToOptimizableRepartitionNodes.get(keyChangingParent)); - keyChangingOperationsToOptimizableRepartitionNodes.remove(keyChangingParent); + mergeNodeKeyChangingParentsToRemove.add(keyChangingParent); } - keyChangingOperationsToOptimizableRepartitionNodes.put(mergeKey, repartitionNodes); } + + for (final StreamsGraphNode mergeNodeKeyChangingParent : mergeNodeKeyChangingParentsToRemove) { + keyChangingOperationsToOptimizableRepartitionNodes.remove(mergeNodeKeyChangingParent); + } } @SuppressWarnings("unchecked") diff --git a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/graph/StreamsGraphTest.java b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/graph/StreamsGraphTest.java index 0fecaa2963eb1..e2006e607e8c8 100644 --- a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/graph/StreamsGraphTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/graph/StreamsGraphTest.java @@ -22,6 +22,7 @@ import org.apache.kafka.streams.StreamsBuilder; import org.apache.kafka.streams.StreamsConfig; import org.apache.kafka.streams.Topology; +import org.apache.kafka.streams.kstream.Consumed; import org.apache.kafka.streams.kstream.Grouped; import org.apache.kafka.streams.kstream.JoinWindows; import org.apache.kafka.streams.kstream.KStream; @@ -124,6 +125,28 @@ public void shouldNotOptimizeWhenAThroughOperationIsDone() { } + @Test + public void shouldOptimizeSeveralMergeNodesWithCommonKeyChangingParent() { + final StreamsBuilder streamsBuilder = new StreamsBuilder(); + final KStream parentStream = streamsBuilder.stream("input_topic", Consumed.with(Serdes.Integer(), Serdes.Integer())) + .selectKey(Integer::sum); + + final KStream childStream1 = parentStream.mapValues(v -> v + 1); + final KStream childStream2 = parentStream.mapValues(v -> v + 2); + final KStream childStream3 = parentStream.mapValues(v -> v + 3); + + childStream1 + .merge(childStream2) + .merge(childStream3) + .to("output_topic"); + + final Properties properties = new Properties(); + properties.setProperty(StreamsConfig.TOPOLOGY_OPTIMIZATION, StreamsConfig.OPTIMIZE); + final Topology topology = streamsBuilder.build(properties); + + assertEquals(expectedMergeOptimizedTopology, topology.describe().toString()); + } + private Topology getTopologyWithChangingValuesAfterChangingKey(final String optimizeConfig) { final StreamsBuilder builder = new StreamsBuilder(); @@ -242,4 +265,30 @@ private int getCountOfRepartitionTopicsFound(final String topologyString) { + " Sink: KSTREAM-SINK-0000000009 (topic: output-topic)\n" + " <-- KSTREAM-MAPVALUES-0000000008\n\n"; + + private String expectedMergeOptimizedTopology = "Topologies:\n" + + " Sub-topology: 0\n" + + " Source: KSTREAM-SOURCE-0000000000 (topics: [input_topic])\n" + + " --> KSTREAM-KEY-SELECT-0000000001\n" + + " Processor: KSTREAM-KEY-SELECT-0000000001 (stores: [])\n" + + " --> KSTREAM-MAPVALUES-0000000002, KSTREAM-MAPVALUES-0000000003, KSTREAM-MAPVALUES-0000000004\n" + + " <-- KSTREAM-SOURCE-0000000000\n" + + " Processor: KSTREAM-MAPVALUES-0000000002 (stores: [])\n" + + " --> KSTREAM-MERGE-0000000005\n" + + " <-- KSTREAM-KEY-SELECT-0000000001\n" + + " Processor: KSTREAM-MAPVALUES-0000000003 (stores: [])\n" + + " --> KSTREAM-MERGE-0000000005\n" + + " <-- KSTREAM-KEY-SELECT-0000000001\n" + + " Processor: KSTREAM-MAPVALUES-0000000004 (stores: [])\n" + + " --> KSTREAM-MERGE-0000000006\n" + + " <-- KSTREAM-KEY-SELECT-0000000001\n" + + " Processor: KSTREAM-MERGE-0000000005 (stores: [])\n" + + " --> KSTREAM-MERGE-0000000006\n" + + " <-- KSTREAM-MAPVALUES-0000000002, KSTREAM-MAPVALUES-0000000003\n" + + " Processor: KSTREAM-MERGE-0000000006 (stores: [])\n" + + " --> KSTREAM-SINK-0000000007\n" + + " <-- KSTREAM-MERGE-0000000005, KSTREAM-MAPVALUES-0000000004\n" + + " Sink: KSTREAM-SINK-0000000007 (topic: output_topic)\n" + + " <-- KSTREAM-MERGE-0000000006\n\n"; + }