[Java] Cache DoFn invoker constructors per instance (#39310)

* Cache DoFn invoker constructors per instance

Co-Authored-By: Codex <noreply@openai.com>

* Add DoFn invoker cache fix to changelog

Co-Authored-By: Codex <noreply@openai.com>

* Clarify DoFn cache performance improvement

Co-Authored-By: Codex <noreply@openai.com>

---------

Co-authored-by: Codex <noreply@openai.com>

Author: lofifnc

CHANGES.md

@@ -83,6 +83,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)).
+* Improved Java pipeline performance by avoiding repeated `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

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<InvokerCacheKey, Constructor<?>> byteBuddyInvokerConstructorCache =
       new ConcurrentHashMap<>();
 
+  /**
+   * A weak, identity-keyed cache of the constructor selected for each {@link DoFn} instance.
+   *
+   * <p>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.
+   *
+   * <p>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<DoFn<?, ?>, Constructor<?>> byteBuddyInvokerConstructorInstanceCache =
+      new MapMaker().weakKeys().makeMap();
+
   private ByteBuddyDoFnInvokerFactory() {}
 
   /** Returns the {@link DoFnInvoker} for the given {@link DoFn}. */
@@ -330,39 +348,15 @@ public <InputT, OutputT> DoFnInvoker<InputT, OutputT> 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<InputT> 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<InputT>) TypeDescriptor.of(Object.class);
-    }
-
-    TypeDescriptor<OutputT> 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<OutputT>) TypeDescriptor.of(Object.class);
-    }
+    Constructor<?> invokerConstructor =
+        byteBuddyInvokerConstructorInstanceCache.computeIfAbsent(
+            fn, unused -> getByteBuddyInvokerConstructor(signature, fn));
 
     try {
       @SuppressWarnings("unchecked")
       DoFnInvokerBase<InputT, OutputT, DoFn<InputT, OutputT>> invoker =
           (DoFnInvokerBase<InputT, OutputT, DoFn<InputT, OutputT>>)
-              getByteBuddyInvokerConstructor(signature, inputType, outputType).newInstance(fn);
+              invokerConstructor.newInstance(fn);
 
       if (signature.onTimerMethods() != null) {
         for (OnTimerMethod onTimerMethod : signature.onTimerMethods().values()) {
@@ -389,6 +383,39 @@ public <InputT, OutputT> DoFnInvoker<InputT, OutputT> newByteBuddyInvoker(
     }
   }
 
+  private <InputT, OutputT> Constructor<?> getByteBuddyInvokerConstructor(
+      DoFnSignature signature, DoFn<InputT, OutputT> 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<InputT> 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<InputT>) TypeDescriptor.of(Object.class);
+    }
+
+    TypeDescriptor<OutputT> 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<OutputT>) 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.

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<T> getInputTypeDescriptor() {
       public TypeDescriptor<T> 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<String, String> stringFn = new DynamicTypeDoFn<>(TypeDescriptors.strings());
@@ -1456,4 +1469,38 @@ public TypeDescriptor<T> getOutputTypeDescriptor() {
         stringInvoker.getClass(),
         intInvoker.getClass());
   }
+
+  @Test
+  public void testTypeDescriptorResolutionIsCachedPerDoFnInstance() {
+    class CountingTypeDoFn extends DoFn<String, String> {
+      private int inputTypeDescriptorCalls;
+      private int outputTypeDescriptorCalls;
+
+      @ProcessElement
+      public void processElement(@Element String element, OutputReceiver<String> out) {
+        out.output(element);
+      }
+
+      @Override
+      public TypeDescriptor<String> getInputTypeDescriptor() {
+        inputTypeDescriptorCalls++;
+        return TypeDescriptors.strings();
+      }
+
+      @Override
+      public TypeDescriptor<String> getOutputTypeDescriptor() {
+        outputTypeDescriptorCalls++;
+        return TypeDescriptors.strings();
+      }
+    }
+
+    CountingTypeDoFn fn = new CountingTypeDoFn();
+
+    DoFnInvoker<String, String> firstInvoker = DoFnInvokers.invokerFor(fn);
+    DoFnInvoker<String, String> secondInvoker = DoFnInvokers.invokerFor(fn);
+
+    assertSame(firstInvoker.getClass(), secondInvoker.getClass());
+    assertEquals(1, fn.inputTypeDescriptorCalls);
+    assertEquals(1, fn.outputTypeDescriptorCalls);
+  }
 }