KAFKA-15045: (KIP-924 pt. 2) Implement ApplicationState and KafkaStreamsState #15920

apourchet · 2024-05-10T21:42:39Z

This PR implements read-only container classes for ApplicationState and KafkaStreamsState, and initializes those within
StreamsPartitionAssignor#assign.

New internal methods were also added to the ClientState to easily pass this data through to the KafkaStreamsState.

One test was added to check the lag sorting within the implementation of KafkaStreamsState, which is the counterpart to the test that existed for the ClientState class.

Committer Checklist (excluded from commit message)

Verify design and implementation
Verify test coverage and CI build status
Verify documentation (including upgrade notes)

This PR implements read-only container classes for ApplicationState and KafkaStreamsState, and initializes those within StreamsPartitionAssignor#assign. New internal methods were also added to the ClientState to easily pass this data through to the KafkaStreamsState.

streams/src/main/java/org/apache/kafka/streams/processor/assignment/ProcessId.java

ableegoldman · 2024-05-11T03:19:06Z

...ams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java

@@ -459,6 +468,38 @@ public GroupAssignment assign(final Cluster metadata, final GroupSubscription gr
        }
    }

+    private ApplicationState getApplicationState(


nit: we don't use "get" in Streams getter names. I guess this isn't exactly a pure getter, but still. On that note, perhaps a better name would be buildApplicationState? 🤔

Also: even though it's all internal, I've been on a crusade to get everyone to write javadocs for methods in the StreamsPartitionAssignor with at least a brief explanation of what it does.

It's just a super complicated class that does a lot and often mutates things in a way that isn't obvious, so every little bit of documentation helps

...ams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java

ableegoldman · 2024-05-11T03:24:16Z

...ams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java

@@ -432,6 +437,10 @@ public GroupAssignment assign(final Cluster metadata, final GroupSubscription gr

            // compute the assignment of tasks to threads within each client and build the final group assignment

+            getApplicationState(


I know we're just throwing away the return value for now but I'd still do this:

Suggested change

getApplicationState(

final ApplicationState applicationState = getApplicationState(

Otherwise it kind of seems like this method is supposed to be mutating the input parameters (many of the StreamsPartitionAssignor methods work this way so it's good to distinguish when we're just building something vs operating on the passed in structures)

SpotBugs was complaining about unused variables when I had that, so I will simply leave the invocation of that new method out entirely.

ah, I see. Well there's no concern about forgetting to put this back in since we'll eventually need the result lol. All good 👍

ableegoldman · 2024-05-11T03:29:41Z

.../main/java/org/apache/kafka/streams/processor/internals/assignment/ApplicationStateImpl.java

+    public ApplicationStateImpl(
+        final AssignmentConfigs assignmentConfigs,
+        final Map<ProcessId, KafkaStreamsState> kafkaStreamsStates,
+        final Set<TaskId> statefulTasks,
+        final Set<TaskId> statelessTasks
+    ) {


KafkaStreams formatting for long signatures is (unfortunately) done like this :

Suggested change

public ApplicationStateImpl(

final AssignmentConfigs assignmentConfigs,

final Map<ProcessId, KafkaStreamsState> kafkaStreamsStates,

final Set<TaskId> statefulTasks,

final Set<TaskId> statelessTasks

) {

public ApplicationStateImpl(final AssignmentConfigs assignmentConfigs,

final Map<ProcessId, KafkaStreamsState> kafkaStreamsStates,

final Set<TaskId> statefulTasks,

final Set<TaskId> statelessTasks) {

It's annoying, I know. Can you make a pass over all the methods and make sure they follow the AK style?

ableegoldman · 2024-05-11T03:31:39Z

.../main/java/org/apache/kafka/streams/processor/internals/assignment/ApplicationStateImpl.java

+        final Set<TaskId> union = new HashSet<>(statefulTasks);
+        union.addAll(statelessTasks);


nit: this could in theory be called multiple times, so we probably want to cache the result instead of building up a new map each time. We can still do it lazily, but I'd say just build the map in the constructor so we can make everything final (and unmodifiable)

ableegoldman · 2024-05-11T03:35:34Z

...main/java/org/apache/kafka/streams/processor/internals/assignment/KafkaStreamsStateImpl.java

+
+    @Override
+    public long lagFor(final TaskId task) {
+        final Long totalLag = taskLagTotals.get(task);


We should also throw if the KafkaStreamsState was built without requesting the task lags be computed, right? (IIRC we decided on UnsupportedOperationException in the last PR)

Same for the other lag-related methods here. I guess the safest thing to do is wrap the taskLagTotals field in an Optional and make it empty when the user passed in computeTaskLags=false to the ApplicationState#kafkaStreamsStates API?

ableegoldman · 2024-05-11T03:38:06Z

...main/java/org/apache/kafka/streams/processor/internals/assignment/KafkaStreamsStateImpl.java

+    public long lagFor(final TaskId task) {
+        final Long totalLag = taskLagTotals.get(task);
+        if (totalLag == null) {
+            throw new IllegalStateException("Tried to lookup lag for unknown task " + task);


As a general rule that we clearly don't always follow in Streams, but should/are actively getting better about, we should make sure to always log an error before throwing a (new) exception. Sometimes the error handling can make it difficult to trace back an exception to the original source of error, and error logs help with that. Doesn't need to be too complicated -- although in this case, perhaps it would make sense to include the keySet of taskLagTotals in the error log (I'm thinking to help differentiate the case where this specific task isn't included vs the entire map being empty)

ableegoldman · 2024-05-11T03:41:08Z

.../main/java/org/apache/kafka/streams/processor/internals/assignment/ApplicationStateImpl.java

+
+    @Override
+    public Map<ProcessId, KafkaStreamsState> kafkaStreamsStates(final boolean computeTaskLags) {
+        return kafkaStreamsStates;


We should make sure to respect the computeTaskLags param here. There's two things we can do: wait until this method and then call some KafkaStreamsStateImpl#computeTaskLags method, or just wait to construct the KafkaStreamsImpls at all until we know whether or not we should compute the task lags, and then pass the computeTaskLags flag into the KafkaStreamsImpl constructor.

I personally prefer the latter since that way there's no partially-initialized classes floating around and we don't have to keep track of when the task lags are computed/initialized.

ableegoldman · 2024-05-11T03:42:36Z

streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java

+        return new TreeSet<>(previousStandbyTasks.taskIds());
+    }
+
+    public SortedMap<String, Set<TaskId>> taskIdsByConsumer() {


nit: call this taskIdsByPreviousConsumer or something to that effect (ie include the word "previous")

It's so hard to keep track of what pertains to the previous assignment vs the new assignment in this mess of a class lol. That's one of the biggest improvements in KIP-924

ableegoldman

LGTM, ping me when the build finishes running and I'll merge

ableegoldman · 2024-05-15T00:29:15Z

The test env seems pretty unstable right now but we have at least one clean build for each java version if you look at the two latest runs. All test failures are unrelated as well. Seems safe to merge

ableegoldman · 2024-05-15T00:30:01Z

Merged to trunk

ableegoldman changed the title ~~KAFKA-15045: (KIP-924): Implement ApplicationState and KafkaStreamsState~~ KAFKA-15045: (KIP-924 pt. 2) Implement ApplicationState and KafkaStreamsState May 11, 2024

ableegoldman reviewed May 11, 2024

View reviewed changes

apourchet added 3 commits May 13, 2024 09:50

addressed first round of comments

e3c16c2

handling lag computation failures properly

715464d

process id string override

19f52be

ableegoldman approved these changes May 13, 2024

View reviewed changes

Empty commit for CI retrigger

ca8f328

ableegoldman merged commit 0c5e8d3 into apache:trunk May 15, 2024
1 check failed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

KAFKA-15045: (KIP-924 pt. 2) Implement ApplicationState and KafkaStreamsState #15920

KAFKA-15045: (KIP-924 pt. 2) Implement ApplicationState and KafkaStreamsState #15920

apourchet commented May 10, 2024

ableegoldman May 11, 2024

ableegoldman May 11, 2024

apourchet May 13, 2024

ableegoldman May 13, 2024

ableegoldman May 11, 2024

ableegoldman May 11, 2024

ableegoldman May 11, 2024

ableegoldman May 11, 2024

ableegoldman May 11, 2024

ableegoldman May 11, 2024

ableegoldman left a comment

ableegoldman commented May 15, 2024

ableegoldman commented May 15, 2024

		@@ -432,6 +437,10 @@ public GroupAssignment assign(final Cluster metadata, final GroupSubscription gr

		// compute the assignment of tasks to threads within each client and build the final group assignment

		getApplicationState(

	getApplicationState(
	final ApplicationState applicationState = getApplicationState(

		final Set<TaskId> union = new HashSet<>(statefulTasks);
		union.addAll(statelessTasks);

KAFKA-15045: (KIP-924 pt. 2) Implement ApplicationState and KafkaStreamsState #15920

KAFKA-15045: (KIP-924 pt. 2) Implement ApplicationState and KafkaStreamsState #15920

Conversation

apourchet commented May 10, 2024

Committer Checklist (excluded from commit message)

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ableegoldman left a comment

Choose a reason for hiding this comment

ableegoldman commented May 15, 2024

ableegoldman commented May 15, 2024