diff --git a/sdks/java/extensions/protobuf/src/main/java/org/apache/beam/sdk/extensions/protobuf/ProtoByteBuddyUtils.java b/sdks/java/extensions/protobuf/src/main/java/org/apache/beam/sdk/extensions/protobuf/ProtoByteBuddyUtils.java index 98f80f6786c8..a1e45868e87d 100644 --- a/sdks/java/extensions/protobuf/src/main/java/org/apache/beam/sdk/extensions/protobuf/ProtoByteBuddyUtils.java +++ b/sdks/java/extensions/protobuf/src/main/java/org/apache/beam/sdk/extensions/protobuf/ProtoByteBuddyUtils.java @@ -18,6 +18,7 @@ package org.apache.beam.sdk.extensions.protobuf; import static org.apache.beam.sdk.extensions.protobuf.ProtoSchemaTranslator.getFieldNumber; +import static org.apache.beam.sdk.util.ByteBuddyUtils.getClassLoadingStrategy; import com.google.protobuf.BoolValue; import com.google.protobuf.ByteString; @@ -55,7 +56,6 @@ import net.bytebuddy.description.type.TypeDescription; import net.bytebuddy.description.type.TypeDescription.ForLoadedType; import net.bytebuddy.dynamic.DynamicType; -import net.bytebuddy.dynamic.loading.ClassLoadingStrategy; import net.bytebuddy.dynamic.scaffold.InstrumentedType; import net.bytebuddy.implementation.FixedValue; import net.bytebuddy.implementation.Implementation; @@ -88,7 +88,6 @@ import org.apache.beam.sdk.schemas.SchemaUserTypeCreator; import org.apache.beam.sdk.schemas.logicaltypes.EnumerationType; import org.apache.beam.sdk.schemas.logicaltypes.OneOfType; -import org.apache.beam.sdk.schemas.utils.ByteBuddyUtils; import org.apache.beam.sdk.schemas.utils.ByteBuddyUtils.ConvertType; import org.apache.beam.sdk.schemas.utils.ByteBuddyUtils.ConvertValueForGetter; import org.apache.beam.sdk.schemas.utils.ByteBuddyUtils.ConvertValueForSetter; @@ -511,8 +510,16 @@ public TypeConversion createSetterConversions(StackManipulati int[] keys = getterMethodMap.keySet().stream().mapToInt(Integer::intValue).toArray(); + Class targetClass = getLoadingTarget(protoClass); + @SuppressWarnings("unchecked") DynamicType.Builder> builder = - ByteBuddyUtils.subclassGetterInterface(BYTE_BUDDY, protoClass, OneOfType.Value.class); + (DynamicType.Builder>) + BYTE_BUDDY + .with(new InjectPackageStrategy(targetClass)) + .subclass( + TypeDescription.Generic.Builder.parameterizedType( + FieldValueGetter.class, protoClass, OneOfType.Value.class) + .build()); builder = builder .method(ElementMatchers.named("name")) @@ -546,7 +553,7 @@ public TypeConversion createSetterConversions(StackManipulati return builder .visit(new AsmVisitorWrapper.ForDeclaredMethods().writerFlags(ClassWriter.COMPUTE_FRAMES)) .make() - .load(ReflectHelpers.findClassLoader(), ClassLoadingStrategy.Default.INJECTION) + .load(targetClass.getClassLoader(), getClassLoadingStrategy(targetClass)) .getLoaded() .getDeclaredConstructor(List.class, OneOfType.class) .newInstance(getters, oneOfType); @@ -568,9 +575,16 @@ FieldValueSetter createOneOfSetter( boolean contiguous = isContiguous(indices); int[] keys = setterMethodMap.keySet().stream().mapToInt(Integer::intValue).toArray(); + Class targetClass = getLoadingTarget(protoBuilderClass); + @SuppressWarnings("unchecked") DynamicType.Builder> builder = - ByteBuddyUtils.subclassSetterInterface( - BYTE_BUDDY, protoBuilderClass, OneOfType.Value.class); + (DynamicType.Builder>) + BYTE_BUDDY + .with(new InjectPackageStrategy(targetClass)) + .subclass( + TypeDescription.Generic.Builder.parameterizedType( + FieldValueSetter.class, protoBuilderClass, OneOfType.Value.class) + .build()); builder = builder .method(ElementMatchers.named("name")) @@ -598,7 +612,7 @@ FieldValueSetter createOneOfSetter( return builder .visit(new AsmVisitorWrapper.ForDeclaredMethods().writerFlags(ClassWriter.COMPUTE_FRAMES)) .make() - .load(ReflectHelpers.findClassLoader(), ClassLoadingStrategy.Default.INJECTION) + .load(targetClass.getClassLoader(), getClassLoadingStrategy(targetClass)) .getLoaded() .getDeclaredConstructor(List.class) .newInstance(setters); @@ -1103,11 +1117,16 @@ static SchemaUserTypeCreator createB List> setters, Schema schema) { try { + Class targetClass = getLoadingTarget(builderClass); + @SuppressWarnings("unchecked") DynamicType.Builder> builder = - (DynamicType.Builder) + (DynamicType.Builder>) BYTE_BUDDY - .with(new InjectPackageStrategy(builderClass)) - .subclass(Supplier.class) + .with(new InjectPackageStrategy(targetClass)) + .subclass( + TypeDescription.Generic.Builder.parameterizedType( + Supplier.class, builderClass) + .build()) .method(ElementMatchers.named("get")) .intercept(new BuilderSupplier(protoClass)); Supplier supplier = @@ -1116,7 +1135,7 @@ static SchemaUserTypeCreator createB new AsmVisitorWrapper.ForDeclaredMethods() .writerFlags(ClassWriter.COMPUTE_FRAMES)) .make() - .load(ReflectHelpers.findClassLoader(), ClassLoadingStrategy.Default.INJECTION) + .load(targetClass.getClassLoader(), getClassLoadingStrategy(targetClass)) .getLoaded() .getDeclaredConstructor() .newInstance(); @@ -1190,4 +1209,12 @@ public ByteCodeAppender appender(final Target implementationTarget) { }; } } + + private static Class getLoadingTarget(Class protoClass) { + ClassLoader loader = ReflectHelpers.findClassLoader(ProtoByteBuddyUtils.class, protoClass); + if (loader == ProtoByteBuddyUtils.class.getClassLoader()) { + return ProtoByteBuddyUtils.class; + } + return protoClass; + } }