bugfix: Avro reflection cache

data-model/src/main/java/org/hypertrace/core/datamodel/shared/AvroBuilderCache.java

@@ -0,0 +1,70 @@
+package org.hypertrace.core.datamodel.shared;
+
+import java.lang.reflect.Constructor;
+import java.lang.reflect.Method;
+import java.util.Arrays;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import javax.annotation.Nonnull;
+import org.apache.avro.specific.SpecificRecord;
+import org.apache.avro.specific.SpecificRecordBuilderBase;
+import org.apache.commons.lang3.tuple.Pair;
+
+/**
+ * Utility class to instantiate a new avro builder in a fast and efficient way. Default Avro
+ * implementation has a perf issue due to uncached reflection based implementation. For more
+ * details, https://issues.apache.org/jira/browse/AVRO-3048
+ *
+ * <p>This avro perf bug is fixed in unreleased avro version (0.11.x). We can get rid of this
+ * utility class once we migrate to Avro 0.11.x or above
+ */
+@SuppressWarnings({"rawtypes", "unchecked"})
+public class AvroBuilderCache {
+  private static final Map<
+          Class<SpecificRecordBuilderBase>, Pair<Method, SpecificRecordBuilderBase>>
+      AVRO_BUILDER_CACHE = new ConcurrentHashMap<>();
+
+  public static <T extends SpecificRecordBuilderBase> T fastNewBuilder(
+      Class<T> specificAvroBuilderClass) {
+    try {
+      Pair<Method, SpecificRecordBuilderBase> pair =
+          AVRO_BUILDER_CACHE.computeIfAbsent(
+              (Class<SpecificRecordBuilderBase>) specificAvroBuilderClass,
+              AvroBuilderCache::newBuilderFor);
+      return (T) pair.getLeft().invoke(null, pair.getRight());
+    } catch (Exception e) {
+      throw new RuntimeException(
+          "Unable to instantiate new builder for: " + specificAvroBuilderClass.getName(), e);
+    }
+  }
+
+  private static Pair<Method, SpecificRecordBuilderBase> newBuilderFor(
+      @Nonnull Class<SpecificRecordBuilderBase> specificAvroBuilderClass) {
+    try {
+      Class<SpecificRecord> specificAvroClass =
+          (Class<SpecificRecord>) specificAvroBuilderClass.getEnclosingClass();
+      Method newBuilderMethod = findBuilderMethod(specificAvroClass);
+
+      final Constructor declaredConstructor = specificAvroBuilderClass.getDeclaredConstructor();
+      declaredConstructor.setAccessible(true);
+      SpecificRecordBuilderBase builderModel =
+          (SpecificRecordBuilderBase) declaredConstructor.newInstance();
+
+      return Pair.of(newBuilderMethod, builderModel);
+    } catch (Exception e) {
+      throw new RuntimeException(
+          "Unable to instantiate the model builder: " + specificAvroBuilderClass.getName(), e);
+    }
+  }
+
+  private static Method findBuilderMethod(Class specificAvroClass) {
+    return Arrays.stream(specificAvroClass.getDeclaredMethods())
+        .filter(method -> method.getName().equals("newBuilder"))
+        .filter(method -> method.getParameterTypes().length == 1)
+        .filter(
+            method ->
+                SpecificRecordBuilderBase.class.isAssignableFrom(method.getParameterTypes()[0]))
+        .findFirst()
+        .orElseThrow(IllegalArgumentException::new);
+  }
+}

data-model/src/main/java/org/hypertrace/core/datamodel/shared/SpanAttributeUtils.java

@@ -1,5 +1,7 @@
 package org.hypertrace.core.datamodel.shared;
 
+import static org.hypertrace.core.datamodel.shared.AvroBuilderCache.fastNewBuilder;
+
 import java.nio.ByteBuffer;
 import java.util.List;
 import java.util.Map;
@@ -34,7 +36,7 @@ public static AttributeValue getAttributeValueWithDefault(
       Event event, String attributeKey, String defaultValue) {
     AttributeValue attributeValue = getAttributeValue(event, attributeKey);
     return attributeValue == null
-        ? AttributeValue.newBuilder().setValue(defaultValue).build()
+        ? fastNewBuilder(AttributeValue.Builder.class).setValue(defaultValue).build()
         : attributeValue;
   }
 

data-model/src/main/java/org/hypertrace/core/datamodel/shared/trace/AttributeValueCreator.java

@@ -1,5 +1,7 @@
 package org.hypertrace.core.datamodel.shared.trace;
 
+import static org.hypertrace.core.datamodel.shared.AvroBuilderCache.fastNewBuilder;
+
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.List;
@@ -10,19 +12,19 @@
 public class AttributeValueCreator {
 
   public static AttributeValue create(String value) {
-    return AttributeValue.newBuilder().setValue(value).build();
+    return fastNewBuilder(AttributeValue.Builder.class).setValue(value).build();
   }
 
   public static AttributeValue create(boolean value) {
-    return AttributeValue.newBuilder().setValue(String.valueOf(value)).build();
+    return fastNewBuilder(AttributeValue.Builder.class).setValue(String.valueOf(value)).build();
   }
 
   public static AttributeValue create(int value) {
-    return AttributeValue.newBuilder().setValue(String.valueOf(value)).build();
+    return fastNewBuilder(AttributeValue.Builder.class).setValue(String.valueOf(value)).build();
   }
 
   public static AttributeValue create(List<String> values) {
-    return AttributeValue.newBuilder().setValueList(values).build();
+    return fastNewBuilder(AttributeValue.Builder.class).setValueList(values).build();
   }
 
   public static AttributeValue createFromByteBuffers(Set<ByteBuffer> values) {
@@ -31,6 +33,6 @@ public static AttributeValue createFromByteBuffers(Set<ByteBuffer> values) {
         value -> {
           list.add(new String(HexUtils.getBytes(value)));
         });
-    return AttributeValue.newBuilder().setValueList(list).build();
+    return fastNewBuilder(AttributeValue.Builder.class).setValueList(list).build();
   }
 }

data-model/src/main/java/org/hypertrace/core/datamodel/shared/trace/MetricValueCreator.java

@@ -1,18 +1,20 @@
 package org.hypertrace.core.datamodel.shared.trace;
 
+import static org.hypertrace.core.datamodel.shared.AvroBuilderCache.fastNewBuilder;
+
 import org.hypertrace.core.datamodel.MetricValue;
 
 public class MetricValueCreator {
 
   public static MetricValue create(double value) {
-    return MetricValue.newBuilder().setValue(value).build();
+    return fastNewBuilder(MetricValue.Builder.class).setValue(value).build();
   }
 
   public static MetricValue create(long value) {
-    return MetricValue.newBuilder().setValue((double) value).build();
+    return fastNewBuilder(MetricValue.Builder.class).setValue((double) value).build();
   }
 
   public static MetricValue create(int value) {
-    return MetricValue.newBuilder().setValue((double) value).build();
+    return fastNewBuilder(MetricValue.Builder.class).setValue((double) value).build();
   }
 }

data-model/src/main/java/org/hypertrace/core/datamodel/shared/trace/StructuredTraceBuilder.java

@@ -1,6 +1,7 @@
 package org.hypertrace.core.datamodel.shared.trace;
 
 import static java.util.Objects.nonNull;
+import static org.hypertrace.core.datamodel.shared.AvroBuilderCache.fastNewBuilder;
 
 import com.google.common.collect.BiMap;
 import com.google.common.collect.HashBiMap;
@@ -182,7 +183,7 @@ public StructuredTrace buildStructuredTrace() {
 
   private StructuredTrace build() {
     // start building the Structured Trace Proto object
-    Builder builder = StructuredTrace.newBuilder();
+    Builder builder = fastNewBuilder(StructuredTrace.Builder.class);
     builder.setCustomerId(customerId);
     builder.setTraceId(traceId);
     builder.setEventList(new ArrayList<>());
@@ -255,7 +256,7 @@ private StructuredTrace build() {
   }
 
   private Edge buildEntityEdge(String parentEntityId, String childEntityId) {
-    Edge.Builder edgeBuilder = Edge.newBuilder();
+    Edge.Builder edgeBuilder = fastNewBuilder(Edge.Builder.class);
     edgeBuilder.setEdgeType(EdgeType.ENTITY_ENTITY);
     edgeBuilder.setSrcIndex(entityIdMapping.get(parentEntityId));
     edgeBuilder.setTgtIndex(entityIdMapping.get(childEntityId));
@@ -276,17 +277,19 @@ private Edge buildEntityEdge(String parentEntityId, String childEntityId) {
   }
 
   private Edge buildEntityEventEdge(String entityId, ByteBuffer eventId) {
-    Edge.Builder edgeBuilder = Edge.newBuilder();
+    Edge.Builder edgeBuilder = fastNewBuilder(Edge.Builder.class);
     edgeBuilder.setEdgeType(EdgeType.ENTITY_EVENT);
     edgeBuilder.setSrcIndex(entityIdMapping.get(entityId));
     edgeBuilder.setTgtIndex(eventIdMapping.get(eventId));
-    edgeBuilder.setAttributes(Attributes.newBuilder().setAttributeMap(new HashMap<>()).build());
-    edgeBuilder.setMetrics(Metrics.newBuilder().setMetricMap(new HashMap<>()).build());
+    edgeBuilder.setAttributes(
+        fastNewBuilder(Attributes.Builder.class).setAttributeMap(new HashMap<>()).build());
+    edgeBuilder.setMetrics(
+        fastNewBuilder(Metrics.Builder.class).setMetricMap(new HashMap<>()).build());
     return edgeBuilder.build();
   }
 
   private Edge buildEventEdge(ByteBuffer parentEventId, ByteBuffer childEventId) {
-    Edge.Builder edgeBuilder = Edge.newBuilder();
+    Edge.Builder edgeBuilder = fastNewBuilder(Edge.Builder.class);
     edgeBuilder.setSrcIndex(eventIdMapping.get(parentEventId));
     edgeBuilder.setTgtIndex(eventIdMapping.get(childEventId));
     edgeBuilder.setEdgeType(EdgeType.EVENT_EVENT);

data-model/src/test/java/org/hypertrace/core/datamodel/shared/AvroBuilderCacheTest.java

@@ -0,0 +1,16 @@
+package org.hypertrace.core.datamodel.shared;
+
+import org.hypertrace.core.datamodel.Metrics;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+public class AvroBuilderCacheTest {
+
+  @Test
+  public void testGetBuilderFromCache() {
+    Metrics.Builder builder1 = AvroBuilderCache.fastNewBuilder(Metrics.Builder.class);
+    Metrics.Builder builder2 = AvroBuilderCache.fastNewBuilder(Metrics.Builder.class);
+    Assertions.assertNotSame(
+        builder1, builder2, "Builders can't be same. Every builder has to be a new instance");
+  }
+}