[KafkaIO] Fix average record size data race and backlog estimation #34165

sjvanrossum · 2025-03-04T17:17:02Z

The offset gap ratio may artificially shrink the backlog if consumers can't catch up to the tail of an expiring topic. This may cause runners to trigger a downscaling event which worsens the issue.

MovingAvg has been modified to atomically write the accumulated state since concurrent normal loads/stores of longs/doubles may tear. The numUpdates field is only used by the writer and can be kept as non-volatile, but the update method ensures that normal loads/stores on numUpdates are ordered in relation to acquiring loads and releasing stores on avg. To prevent false sharing I've padded the class since there may be tens to hundreds of instances of the accumulator and updates happen per consumed record.

The JMH benchmark I've added shows a slight uplift in average time per op for both reads and writes compared to the current implementation.

Results of task :sdks:java:io:kafka:jmh:jmh on a t2d-standard-60 Cloud Workstation:

Benchmark                                 Mode  Cnt       Score       Error  Units
KafkaIOUtilsBenchmark.Atomic              avgt   15   50693.751 ±  4937.770  ns/op
KafkaIOUtilsBenchmark.Atomic:atomicRead   avgt   15    5577.357 ±  1135.962  ns/op
KafkaIOUtilsBenchmark.Atomic:atomicWrite  avgt   15  140926.539 ± 14847.095  ns/op
KafkaIOUtilsBenchmark.Plain               avgt   15   65018.754 ±  9814.457  ns/op
KafkaIOUtilsBenchmark.Plain:plainRead     avgt   15    6658.736 ±   288.912  ns/op
KafkaIOUtilsBenchmark.Plain:plainWrite    avgt   15  181738.789 ± 29403.883  ns/op

Note that this test likely does not highlight the effect of padding since it doesn't construct a large pool of accumulators.

Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:

Mention the appropriate issue in your description (for example: addresses #123), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, comment fixes #<ISSUE NUMBER> instead.
Update CHANGES.md with noteworthy changes.
If this contribution is large, please file an Apache Individual Contributor License Agreement.

See the Contributor Guide for more tips on how to make review process smoother.

To check the build health, please visit https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md

GitHub Actions Tests Status (on master branch)

See CI.md for more information about GitHub Actions CI or the workflows README to see a list of phrases to trigger workflows.

github-actions · 2025-03-04T19:06:01Z

Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment assign set of reviewers

sjvanrossum · 2025-03-05T14:27:55Z

Run Spotless PreCommit

github-actions · 2025-03-05T16:08:43Z

Assigning reviewers. If you would like to opt out of this review, comment assign to next reviewer:

R: @kennknowles for label java.
R: @Abacn for label build.
R: @Abacn for label io.
R: @johnjcasey for label kafka.

Available commands:

stop reviewer notifications - opt out of the automated review tooling
remind me after tests pass - tag the comment author after tests pass
waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

kennknowles

Happy for the improvement. I may be misunderstanding what is strictly necessary to make your changes work as intended but TL;DR the inheritance all seems extraneous - one of them seems like inlining is equivalent and clearer, while the other seems like it is more clearly expressed as a field.

kennknowles · 2025-03-05T16:58:57Z

sdks/java/io/kafka/jmh/build.gradle

+ * limitations under the License.
+ */
+
+plugins { id 'org.apache.beam.module' }


Please use build.gradle.kts for new files

kennknowles · 2025-03-07T15:48:08Z

sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIOUtils.java

+   */
+  @SuppressFBWarnings("UUF_UNUSED_FIELD")
+  private static class MovingAvgPadding {
+    byte p000, p001, p002, p003, p004, p005, p006, p007;


Was this a reaction to a measured slowdown? I'm not like super opposed to having something like this, but I expect it's relevance and effectiveness to wane as the codebase evolves, unless we have some fairly deep and clear understanding of exactly what needs to be maintained.

kennknowles · 2025-03-07T15:48:33Z

sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIOUtils.java

+
+  // The accumulator's fields should be padded to at least 128 bytes (at least 1 or 2
+  // cache lines).
+  private static class MovingAvgFields extends MovingAvgPadding {


Why is the padding not just inlined here? It would be clearer what is going on and saves a layer of inheritance.

kennknowles · 2025-03-07T15:51:04Z

sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIOUtils.java

+   * Sanity of visibility is only useful when the writer thread changes since avg is the only field that can be shared between multiple concurrent threads.
+   */
+  @SuppressFBWarnings("UUF_UNUSED_FIELD")
+  public static final class MovingAvg extends MovingAvgFields {


This class should simply have a private MovingAvgFields member.

kennknowles · 2025-03-07T15:52:50Z

sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIOUtils.java

+      return Double.longBitsToDouble(avg);
+    }
+
+    protected void setAvg(final double value) {


Just make these package-private and call them from the MovingAvg class. (see below: use a field instead of inheritance)

github-actions bot added java build io kafka labels Mar 4, 2025

github-actions bot added the Next Action: Reviewers label Mar 5, 2025

sjvanrossum added 7 commits March 6, 2025 22:30

Fix data race in MovingAvg, remove offset gap ratio

1a2719a

Suppress UUF_UNUSED_FIELD warning

e64deb3

Add package-info.java

4dbaa22

Fix sloppy copy/paste

f13cf61

Move JMH content to jmh package

f4194cb

Change MovingAvg's class and method visibility

3ffeb79

Remove comment

ab5b88d

sjvanrossum force-pushed the kafkaio-sdf-backlog-estimation branch from 52e2c5f to ab5b88d Compare March 6, 2025 22:33

sjvanrossum mentioned this pull request Mar 7, 2025

[KafkaIO] Remove duplicate offset in range check #34201

Open

3 tasks

kennknowles requested changes Mar 7, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[KafkaIO] Fix average record size data race and backlog estimation #34165

[KafkaIO] Fix average record size data race and backlog estimation #34165

sjvanrossum commented Mar 4, 2025 •

edited

Loading

github-actions bot commented Mar 4, 2025

sjvanrossum commented Mar 5, 2025

github-actions bot commented Mar 5, 2025

kennknowles left a comment

kennknowles Mar 5, 2025

kennknowles Mar 7, 2025

kennknowles Mar 7, 2025

kennknowles Mar 7, 2025

kennknowles Mar 7, 2025

[KafkaIO] Fix average record size data race and backlog estimation #34165

Are you sure you want to change the base?

[KafkaIO] Fix average record size data race and backlog estimation #34165

Conversation

sjvanrossum commented Mar 4, 2025 • edited Loading

GitHub Actions Tests Status (on master branch)

github-actions bot commented Mar 4, 2025

sjvanrossum commented Mar 5, 2025

github-actions bot commented Mar 5, 2025

kennknowles left a comment

Choose a reason for hiding this comment

kennknowles Mar 5, 2025

Choose a reason for hiding this comment

kennknowles Mar 7, 2025

Choose a reason for hiding this comment

kennknowles Mar 7, 2025

Choose a reason for hiding this comment

kennknowles Mar 7, 2025

Choose a reason for hiding this comment

kennknowles Mar 7, 2025

Choose a reason for hiding this comment

sjvanrossum commented Mar 4, 2025 •

edited

Loading