From 037b40822a87e909d0438bfbf7852e858b4a8f3f Mon Sep 17 00:00:00 2001 From: Alex Kolb Date: Mon, 13 Jul 2026 22:11:12 +0200 Subject: [PATCH 1/2] Cache DoFn invoker constructors per instance Co-Authored-By: Codex --- .../reflect/ByteBuddyDoFnInvokerFactory.java | 83 ++++++++++++------- .../transforms/reflect/DoFnInvokersTest.java | 47 +++++++++++ 2 files changed, 102 insertions(+), 28 deletions(-) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyDoFnInvokerFactory.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyDoFnInvokerFactory.java index c08243fda5c0..14f864c5b042 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyDoFnInvokerFactory.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyDoFnInvokerFactory.java @@ -29,6 +29,7 @@ import java.util.Map; import java.util.Objects; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; import net.bytebuddy.ByteBuddy; import net.bytebuddy.description.field.FieldDescription; import net.bytebuddy.description.method.MethodDescription; @@ -113,6 +114,7 @@ import org.apache.beam.sdk.values.TypeDescriptor; import org.apache.beam.sdk.values.TypeDescriptors; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.MapMaker; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.primitives.Primitives; import org.checkerframework.checker.nullness.qual.MonotonicNonNull; @@ -242,6 +244,22 @@ public String toString() { private final Map> byteBuddyInvokerConstructorCache = new ConcurrentHashMap<>(); + /** + * A weak, identity-keyed cache of the constructor selected for each {@link DoFn} instance. + * + *

Selecting a constructor resolves the {@link DoFn}'s input and output type descriptors. Some + * runners create a new invoker for the same {@link DoFn} at every bundle boundary, so doing that + * work on every invocation can be expensive. A deserialized {@link DoFn} is a new instance, so + * its constructor is selected again; later invoker creation for that instance reuses the cached + * selection. + * + *

The constructor values do not reference their {@link DoFn} keys, and weak keys allow unused + * instances to be collected. {@link MapMaker#weakKeys} also uses identity equality, which is + * required because separate instances of the same class may have different generic types. + */ + private final ConcurrentMap, Constructor> byteBuddyInvokerConstructorInstanceCache = + new MapMaker().weakKeys().makeMap(); + private ByteBuddyDoFnInvokerFactory() {} /** Returns the {@link DoFnInvoker} for the given {@link DoFn}. */ @@ -330,39 +348,15 @@ public DoFnInvoker newByteBuddyInvoker( signature.fnClass(), fn.getClass()); - // Extract input and output type descriptors to distinguish generic instantiations. - // Fall back to Object.class if unavailable. When type info is lost, different generic - // instantiations share an invoker, which is acceptable since the DoFn class in the cache - // key prevents collisions between different DoFn classes. - TypeDescriptor inputType; - try { - inputType = fn.getInputTypeDescriptor(); - } catch (Exception e) { - // Some DoFns (like MapElements) throw IllegalStateException if queried after - // serialization. - // In this case, we fall back to the raw class behavior (Object). - inputType = null; - } - if (inputType == null) { - inputType = (TypeDescriptor) TypeDescriptor.of(Object.class); - } - - TypeDescriptor outputType; - try { - outputType = fn.getOutputTypeDescriptor(); - } catch (Exception e) { - // Same as above: fall back to Object if type info is unavailable. - outputType = null; - } - if (outputType == null) { - outputType = (TypeDescriptor) TypeDescriptor.of(Object.class); - } + Constructor invokerConstructor = + byteBuddyInvokerConstructorInstanceCache.computeIfAbsent( + fn, unused -> getByteBuddyInvokerConstructor(signature, fn)); try { @SuppressWarnings("unchecked") DoFnInvokerBase> invoker = (DoFnInvokerBase>) - getByteBuddyInvokerConstructor(signature, inputType, outputType).newInstance(fn); + invokerConstructor.newInstance(fn); if (signature.onTimerMethods() != null) { for (OnTimerMethod onTimerMethod : signature.onTimerMethods().values()) { @@ -389,6 +383,39 @@ public DoFnInvoker newByteBuddyInvoker( } } + private Constructor getByteBuddyInvokerConstructor( + DoFnSignature signature, DoFn fn) { + // Extract input and output type descriptors to distinguish generic instantiations. + // Fall back to Object.class if unavailable. When type info is lost, different generic + // instantiations share an invoker, which is acceptable since the DoFn class in the cache + // key prevents collisions between different DoFn classes. + TypeDescriptor inputType; + try { + inputType = fn.getInputTypeDescriptor(); + } catch (Exception e) { + // Some DoFns (like MapElements) throw IllegalStateException if queried after + // serialization. + // In this case, we fall back to the raw class behavior (Object). + inputType = null; + } + if (inputType == null) { + inputType = (TypeDescriptor) TypeDescriptor.of(Object.class); + } + + TypeDescriptor outputType; + try { + outputType = fn.getOutputTypeDescriptor(); + } catch (Exception e) { + // Same as above: fall back to Object if type info is unavailable. + outputType = null; + } + if (outputType == null) { + outputType = (TypeDescriptor) TypeDescriptor.of(Object.class); + } + + return getByteBuddyInvokerConstructor(signature, inputType, outputType); + } + /** * Returns a generated constructor for a {@link DoFnInvoker} for the given {@link DoFnSignature} * and specific generic types. diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnInvokersTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnInvokersTest.java index e9269deb18ea..9f4fdb6c9c28 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnInvokersTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnInvokersTest.java @@ -82,6 +82,7 @@ import org.apache.beam.sdk.values.TypeDescriptor; import org.apache.beam.sdk.values.TypeDescriptors; import org.apache.beam.sdk.values.WindowedValues; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; import org.junit.Before; import org.junit.Rule; @@ -1440,6 +1441,18 @@ public TypeDescriptor getInputTypeDescriptor() { public TypeDescriptor getOutputTypeDescriptor() { return typeDescriptor; } + + // Deliberately make separate instances equal. The per-instance cache must use identity + // because equal instances can have different generic types. + @Override + public boolean equals(@Nullable Object other) { + return other instanceof DynamicTypeDoFn; + } + + @Override + public int hashCode() { + return 0; + } } DoFn stringFn = new DynamicTypeDoFn<>(TypeDescriptors.strings()); @@ -1456,4 +1469,38 @@ public TypeDescriptor getOutputTypeDescriptor() { stringInvoker.getClass(), intInvoker.getClass()); } + + @Test + public void testTypeDescriptorResolutionIsCachedPerDoFnInstance() { + class CountingTypeDoFn extends DoFn { + private int inputTypeDescriptorCalls; + private int outputTypeDescriptorCalls; + + @ProcessElement + public void processElement(@Element String element, OutputReceiver out) { + out.output(element); + } + + @Override + public TypeDescriptor getInputTypeDescriptor() { + inputTypeDescriptorCalls++; + return TypeDescriptors.strings(); + } + + @Override + public TypeDescriptor getOutputTypeDescriptor() { + outputTypeDescriptorCalls++; + return TypeDescriptors.strings(); + } + } + + CountingTypeDoFn fn = new CountingTypeDoFn(); + + DoFnInvoker firstInvoker = DoFnInvokers.invokerFor(fn); + DoFnInvoker secondInvoker = DoFnInvokers.invokerFor(fn); + + assertSame(firstInvoker.getClass(), secondInvoker.getClass()); + assertEquals(1, fn.inputTypeDescriptorCalls); + assertEquals(1, fn.outputTypeDescriptorCalls); + } } From 5e0cc1a2f6fb31e720cdded3dbb92dfb6ee0574b Mon Sep 17 00:00:00 2001 From: Alex Kolb Date: Mon, 13 Jul 2026 22:22:28 +0200 Subject: [PATCH 2/2] Add DoFn invoker cache fix to changelog Co-Authored-By: Codex --- CHANGES.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGES.md b/CHANGES.md index 1e09ade6c389..7a465ba8011e 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -82,6 +82,7 @@ ## Bugfixes * Fixed unbounded checkpoint state growth for splittable DoFns that self-checkpoint on the portable Flink runner (Java) ([#27648](https://github.com/apache/beam/issues/27648)). +* Avoided repeated Java `DoFn` type descriptor resolution when creating cached invokers ([#39309](https://github.com/apache/beam/issues/39309)). * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). ## Security Fixes