Skip to content

Commit 4d337b6

Browse files
committed
[Java] Fix DoFnInvoker cache collision for generic types (with fallback)
This PR fixes a critical issue where ByteBuddyDoFnInvokerFactory failed to distinguish between different generic instantiations of the same DoFn class (e.g., MyFn<String> vs MyFn<Integer>). 1. Cache Key Strategy: Introduced InvokerCacheKey to include input/output TypeDescriptors in the cache lookup. 2. Class Naming: Updated generateInvokerClass to append a type-based hash suffix to ensure unique class names. 3. Robustness (The Fix): Added defensive try-catch blocks when accessing TypeDescriptors. - Some internal transforms (like MapElements) throw IllegalStateException if getOutputTypeDescriptor() is called after serialization. - In these cases, the factory now gracefully falls back to using Object.class (legacy behavior), ensuring backward compatibility for transforms that do not retain type information at runtime. Fixes #37351
1 parent 4f60c59 commit 4d337b6

1 file changed

Lines changed: 19 additions & 5 deletions

File tree

sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyDoFnInvokerFactory.java

Lines changed: 19 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -318,14 +318,28 @@ public <InputT, OutputT> DoFnInvoker<InputT, OutputT> newByteBuddyInvoker(
318318
fn.getClass());
319319

320320
// Extract input and output type descriptors from the DoFn instance
321-
// Fall back to Object.class if the type descriptors are null (e.g., for mocked DoFn instances)
322-
@SuppressWarnings("unchecked")
323-
TypeDescriptor<InputT> inputType = fn.getInputTypeDescriptor();
321+
// Fall back to Object.class if the type descriptors are null or unavailable (e.g., MapElements
322+
// after serialization)
323+
TypeDescriptor<InputT> inputType;
324+
try {
325+
inputType = fn.getInputTypeDescriptor();
326+
} catch (Exception e) {
327+
// Some DoFns (like MapElements) throw IllegalStateException if queried after
328+
// serialization.
329+
// In this case, we fall back to the raw class behavior (Object).
330+
inputType = null;
331+
}
324332
if (inputType == null) {
325333
inputType = (TypeDescriptor<InputT>) TypeDescriptor.of(Object.class);
326334
}
327-
@SuppressWarnings("unchecked")
328-
TypeDescriptor<OutputT> outputType = fn.getOutputTypeDescriptor();
335+
336+
TypeDescriptor<OutputT> outputType;
337+
try {
338+
outputType = fn.getOutputTypeDescriptor();
339+
} catch (Exception e) {
340+
// Same as above: fall back to Object if type info is unavailable.
341+
outputType = null;
342+
}
329343
if (outputType == null) {
330344
outputType = (TypeDescriptor<OutputT>) TypeDescriptor.of(Object.class);
331345
}

0 commit comments

Comments
 (0)