sebastienviale commented on code in PR #22725:
URL: https://github.com/apache/kafka/pull/22725#discussion_r3614850799
##########
streams/test-utils/src/main/java/org/apache/kafka/streams/TopologyTestDriver.java:
##########
@@ -979,7 +1118,91 @@ private StateStore getStateStore(final String name,
return null;
}
- private void throwIfBuiltInStore(final StateStore stateStore) {
+ /**
+ * Return the {@link StateStore} for the task owning {@code partition} of
the sub-topology that
+ * registers a store named {@code name}. If the store name appears in
multiple
+ * sub-topologies, throws {@link IllegalStateException}.
+ *
+ * @param name the store name
+ * @param partition the partition whose owning task should be queried
+ * @return the {@link StateStore}, or {@code null} if no sub-topology
registers a store with this name
+ */
+ public StateStore getStateStore(final String name, final int partition) {
+ requireMultiPartitionMode();
+ return runtime.getStateStore(name, partition);
+ }
+
+ /**
+ * Guard for the partition-aware accessors below ({@link
#getStateStore(String, int)} and
+ * friends). Unlike the original implementation, this never activates
multi-partition mode as a
+ * side effect of what looks like a read-only getter: a {@code getXxx()}
method that can silently
+ * rebuild the entire task graph on first call -- and behave differently
the second time it's
+ * called -- is surprising and hides a non-trivial effect behind an
innocuous-looking signature.
+ *
+ * <p>Multi-partition mode must already be active by the time these
accessors are called, either
+ * because {@link TopologyTestDriverBuilder#build()} activated it (at
least one declared topic has
+ * more than one partition), or because the caller invoked {@link
#activateMultiPartitionMode()}
+ * explicitly. If it isn't, this throws -- it does not activate it for you.
+ *
+ * @throws IllegalStateException if the driver is not operating in
multi-partition mode
+ */
+ private void requireMultiPartitionMode() {
+ if (!multiPartitionModeActive) {
+ throw new IllegalStateException(
+ "This driver is not operating in multi-partition mode. Declare
a topic with more than "
+ + "one partition (via
TopologyTestDriverBuilder#declareTopic() or declareTopic()) "
Review Comment:
I simplified the logic. There is now only one way to activate
multi-partition mode: through `TopologyTestDriverBuilder#declareTopic()` with a
topic having more than one partition.
##########
streams/test-utils/src/main/java/org/apache/kafka/streams/TopologyTestDriver.java:
##########
@@ -812,6 +849,105 @@ public final <K, V> TestOutputTopic<K, V>
createOutputTopic(final String topicNa
return new TestOutputTopic<>(this, topicName, keyDeserializer,
valueDeserializer);
}
+ /**
+ * Declare the number of partitions for an input, output, or generated
repartition topic.
+ * Must be called before any record is piped. Subsequent calls with the
same count are no-ops; calls
+ * with a different count throw {@link IllegalArgumentException}. Calls
after the driver has been
+ * initialised throw {@link IllegalStateException}.
+ *
+ * @param topicName the topic to declare
+ * @param partitions the number of partitions (must be at least 1)
+ * @throws IllegalStateException if the driver has already been
initialised
+ * @throws IllegalArgumentException if {@code partitions} is less than 1,
or the topic was already declared with a different count
+ */
+ void declareTopic(final String topicName, final int partitions) {
+ Objects.requireNonNull(topicName, "topicName cannot be null");
+ if (multiPartitionModeActive) {
+ throw new IllegalStateException(
+ "Cannot declare topic '" + topicName + "' after
multi-partition mode has been activated; "
+ + "declare all multi-partition topics before
piping records.");
+ }
+ if (partitions < 1) {
+ throw new IllegalArgumentException(
+ "Partition count must be at least 1 (topic='" + topicName
+ "', partitions=" + partitions + ").");
+ }
+ final Integer existing = declaredPartitionsByTopic.get(topicName);
+ if (existing != null && existing != partitions) {
+ throw new IllegalArgumentException(
+ "Topic '" + topicName + "' was already declared with " +
existing
+ + " partitions; cannot redeclare with " +
partitions + ".");
+ }
+ declaredPartitionsByTopic.put(topicName, partitions);
+ }
+
+ /**
+ * Activate multi-partition mode. Idempotent. Call this after declaring
all multi-partition
+ * topics and before piping records. The single-partition back-compat path
auto-activates on first use,
+ * so existing tests do not need to call this method.
+ *
+ * <p>This builds the sub-topology task graph: for each sub-topology, it
constructs its
+ * {@link ProcessorTopology}, resolves the partition count of any internal
repartition topic
+ * (declared explicit count > co-partition group inheritance > max
upstream sources >
+ * fallback to 1), validates co-partitioning, and computes the
per-sub-topology partition count
+ * as the max across its source topics.</p>
+ */
+ void activateMultiPartitionMode() {
+ if (multiPartitionModeActive) {
+ return;
Review Comment:
Good point. Since this is a package-private internal method, it's probably
better to throw an IllegalStateException if it's invoked more than once.
##########
streams/test-utils/src/main/java/org/apache/kafka/streams/TopologyTestDriver.java:
##########
@@ -812,6 +849,105 @@ public final <K, V> TestOutputTopic<K, V>
createOutputTopic(final String topicNa
return new TestOutputTopic<>(this, topicName, keyDeserializer,
valueDeserializer);
}
+ /**
+ * Declare the number of partitions for an input, output, or generated
repartition topic.
+ * Must be called before any record is piped. Subsequent calls with the
same count are no-ops; calls
+ * with a different count throw {@link IllegalArgumentException}. Calls
after the driver has been
+ * initialised throw {@link IllegalStateException}.
+ *
+ * @param topicName the topic to declare
+ * @param partitions the number of partitions (must be at least 1)
+ * @throws IllegalStateException if the driver has already been
initialised
+ * @throws IllegalArgumentException if {@code partitions} is less than 1,
or the topic was already declared with a different count
+ */
+ void declareTopic(final String topicName, final int partitions) {
+ Objects.requireNonNull(topicName, "topicName cannot be null");
+ if (multiPartitionModeActive) {
+ throw new IllegalStateException(
+ "Cannot declare topic '" + topicName + "' after
multi-partition mode has been activated; "
+ + "declare all multi-partition topics before
piping records.");
+ }
+ if (partitions < 1) {
+ throw new IllegalArgumentException(
+ "Partition count must be at least 1 (topic='" + topicName
+ "', partitions=" + partitions + ").");
+ }
+ final Integer existing = declaredPartitionsByTopic.get(topicName);
+ if (existing != null && existing != partitions) {
+ throw new IllegalArgumentException(
+ "Topic '" + topicName + "' was already declared with " +
existing
+ + " partitions; cannot redeclare with " +
partitions + ".");
+ }
+ declaredPartitionsByTopic.put(topicName, partitions);
+ }
+
+ /**
+ * Activate multi-partition mode. Idempotent. Call this after declaring
all multi-partition
+ * topics and before piping records. The single-partition back-compat path
auto-activates on first use,
+ * so existing tests do not need to call this method.
Review Comment:
done
##########
streams/test-utils/src/main/java/org/apache/kafka/streams/TopologyTestDriver.java:
##########
@@ -812,6 +849,105 @@ public final <K, V> TestOutputTopic<K, V>
createOutputTopic(final String topicNa
return new TestOutputTopic<>(this, topicName, keyDeserializer,
valueDeserializer);
}
+ /**
+ * Declare the number of partitions for an input, output, or generated
repartition topic.
+ * Must be called before any record is piped. Subsequent calls with the
same count are no-ops; calls
+ * with a different count throw {@link IllegalArgumentException}. Calls
after the driver has been
+ * initialised throw {@link IllegalStateException}.
Review Comment:
Indeed, I removed the the IllegalArgumentExceptions
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
