Add weak cache of built parsers across MessageMarshaller instances

src/main/java/org/curioswitch/common/protobuf/json/MarshallerCache.java

@@ -0,0 +1,68 @@
+/*
+ * Copyright (c) Choko ([email protected])
+ * SPDX-License-Identifier: MIT
+ */
+
+package org.curioswitch.common.protobuf.json;
+
+import java.lang.ref.ReferenceQueue;
+import java.lang.ref.WeakReference;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.locks.ReentrantLock;
+import javax.annotation.Nullable;
+
+// Because it is fairly expensive to build a TypeSpecificMarshaller, we go ahead
+// and cache to save time across different MessageMarshaller instances. We still
+// want to make sure they can be garbage collected, so we use weak references.
+final class MarshallerCache {
+
+  private final Map<MarshallerOptions, MarshallerReference> cache = new HashMap<>();
+  private final ReferenceQueue<TypeSpecificMarshaller<?>> queue = new ReferenceQueue<>();
+  private final ReentrantLock lock = new ReentrantLock();
+
+  @Nullable
+  TypeSpecificMarshaller<?> get(MarshallerOptions key) {
+    lock.lock();
+    try {
+      clean();
+      MarshallerReference ref = cache.get(key);
+      return ref != null ? ref.get() : null;
+    } finally {
+      lock.unlock();
+    }
+  }
+
+  void put(MarshallerOptions key, TypeSpecificMarshaller<?> value) {
+    lock.lock();
+    try {
+      clean();
+      cache.put(key, new MarshallerReference(key, value, queue));
+    } finally {
+      lock.unlock();
+    }
+  }
+
+  private void clean() {
+    MarshallerReference ref;
+    while ((ref = (MarshallerReference) queue.poll()) != null) {
+      cache.remove(ref.getKey());
+    }
+  }
+
+  private static class MarshallerReference extends WeakReference<TypeSpecificMarshaller<?>> {
+    private final MarshallerOptions key;
+
+    MarshallerReference(
+        MarshallerOptions key,
+        TypeSpecificMarshaller<?> value,
+        ReferenceQueue<TypeSpecificMarshaller<?>> queue) {
+      super(value, queue);
+      this.key = key;
+    }
+
+    MarshallerOptions getKey() {
+      return key;
+    }
+  }
+}

src/main/java/org/curioswitch/common/protobuf/json/MarshallerOptions.java

@@ -0,0 +1,97 @@
+/*
+ * Copyright (c) Choko ([email protected])
+ * SPDX-License-Identifier: MIT
+ */
+
+package org.curioswitch.common.protobuf.json;
+
+import com.google.protobuf.Descriptors.FieldDescriptor;
+import com.google.protobuf.Message;
+import java.util.Objects;
+import java.util.Set;
+
+class MarshallerOptions {
+
+  private final Message prototype;
+  private final boolean includingDefaultValueFields;
+  private final Set<FieldDescriptor> fieldsToAlwaysOutput;
+  private final boolean preservingProtoFieldNames;
+  private final boolean ignoringUnknownFields;
+  private final boolean printingEnumsAsInts;
+  private final boolean sortingMapKeys;
+
+  MarshallerOptions(
+      Message prototype,
+      boolean includingDefaultValueFields,
+      Set<FieldDescriptor> fieldsToAlwaysOutput,
+      boolean preservingProtoFieldNames,
+      boolean ignoringUnknownFields,
+      boolean printingEnumsAsInts,
+      boolean sortingMapKeys) {
+    this.prototype = prototype;
+    this.includingDefaultValueFields = includingDefaultValueFields;
+    this.fieldsToAlwaysOutput = fieldsToAlwaysOutput;
+    this.preservingProtoFieldNames = preservingProtoFieldNames;
+    this.ignoringUnknownFields = ignoringUnknownFields;
+    this.printingEnumsAsInts = printingEnumsAsInts;
+    this.sortingMapKeys = sortingMapKeys;
+  }
+
+  Message getPrototype() {
+    return prototype;
+  }
+
+  public boolean isIncludingDefaultValueFields() {
+    return includingDefaultValueFields;
+  }
+
+  Set<FieldDescriptor> getFieldsToAlwaysOutput() {
+    return fieldsToAlwaysOutput;
+  }
+
+  boolean isPreservingProtoFieldNames() {
+    return preservingProtoFieldNames;
+  }
+
+  boolean isIgnoringUnknownFields() {
+    return ignoringUnknownFields;
+  }
+
+  boolean isPrintingEnumsAsInts() {
+    return printingEnumsAsInts;
+  }
+
+  boolean isSortingMapKeys() {
+    return sortingMapKeys;
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) {
+      return true;
+    }
+    if (!(o instanceof MarshallerOptions)) {
+      return false;
+    }
+    MarshallerOptions that = (MarshallerOptions) o;
+    return includingDefaultValueFields == that.includingDefaultValueFields
+        && preservingProtoFieldNames == that.preservingProtoFieldNames
+        && ignoringUnknownFields == that.ignoringUnknownFields
+        && printingEnumsAsInts == that.printingEnumsAsInts
+        && sortingMapKeys == that.sortingMapKeys
+        && prototype.getDescriptorForType().equals(that.prototype.getDescriptorForType())
+        && fieldsToAlwaysOutput.equals(that.fieldsToAlwaysOutput);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(
+        prototype.getDescriptorForType(),
+        includingDefaultValueFields,
+        fieldsToAlwaysOutput,
+        preservingProtoFieldNames,
+        ignoringUnknownFields,
+        printingEnumsAsInts,
+        sortingMapKeys);
+  }
+}

src/main/java/org/curioswitch/common/protobuf/json/MessageMarshaller.java

@@ -492,13 +492,14 @@ public MessageMarshaller build() {
 
       for (Message prototype : prototypes) {
         TypeSpecificMarshaller.buildAndAdd(
-            prototype,
-            includingDefaultValueFields,
-            fieldsToAlwaysOutput,
-            preservingProtoFieldNames,
-            ignoringUnknownFields,
-            printingEnumsAsInts,
-            sortingMapKeys,
+            new MarshallerOptions(
+                prototype,
+                includingDefaultValueFields,
+                fieldsToAlwaysOutput,
+                preservingProtoFieldNames,
+                ignoringUnknownFields,
+                printingEnumsAsInts,
+                sortingMapKeys),
             builtParsers);
       }
 

src/main/java/org/curioswitch/common/protobuf/json/TypeSpecificMarshaller.java

@@ -24,7 +24,6 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 import net.bytebuddy.ByteBuddy;
 import net.bytebuddy.asm.AsmVisitorWrapper.ForDeclaredMethods;
 import net.bytebuddy.description.type.TypeDefinition;
@@ -39,6 +38,8 @@
  */
 public abstract class TypeSpecificMarshaller<T extends Message> {
 
+  private static final MarshallerCache MARSHALLER_CACHE = new MarshallerCache();
+
   private final T prototype;
 
   protected TypeSpecificMarshaller(T prototype) {
@@ -131,27 +132,12 @@ T getMarshalledPrototype() {
     return prototype;
   }
 
-  static <T extends Message> void buildAndAdd(
-      T prototype,
-      boolean includingDefaultValueFields,
-      Set<FieldDescriptor> fieldsToAlwaysOutput,
-      boolean preservingProtoFieldNames,
-      boolean ignoringUnknownFields,
-      boolean printingEnumsAsInts,
-      boolean sortingMapKeys,
-      Map<Descriptor, TypeSpecificMarshaller<?>> builtMarshallers) {
-    if (builtMarshallers.containsKey(prototype.getDescriptorForType())) {
+  static void buildAndAdd(
+      MarshallerOptions options, Map<Descriptor, TypeSpecificMarshaller<?>> builtMarshallers) {
+    if (builtMarshallers.containsKey(options.getPrototype().getDescriptorForType())) {
       return;
     }
-    buildOrFindMarshaller(
-        prototype,
-        includingDefaultValueFields,
-        fieldsToAlwaysOutput,
-        preservingProtoFieldNames,
-        ignoringUnknownFields,
-        printingEnumsAsInts,
-        sortingMapKeys,
-        builtMarshallers);
+    buildOrFindMarshaller(options, builtMarshallers);
     Map<String, TypeSpecificMarshaller<?>> builtMarshallersByFieldName = new HashMap<>();
     for (Map.Entry<Descriptor, TypeSpecificMarshaller<?>> entry : builtMarshallers.entrySet()) {
       builtMarshallersByFieldName.put(
@@ -178,22 +164,16 @@ static <T extends Message> void buildAndAdd(
   }
 
   private static <T extends Message> void buildOrFindMarshaller(
-      T prototype,
-      boolean includingDefaultValueFields,
-      Set<FieldDescriptor> fieldsToAlwaysOutput,
-      boolean preservingProtoFieldNames,
-      boolean ignoringUnknownFields,
-      boolean printingEnumsAsInts,
-      boolean sortingMapKeys,
+      MarshallerOptions options,
       Map<Descriptor, TypeSpecificMarshaller<?>> alreadyBuiltMarshallers) {
-    Descriptor descriptor = prototype.getDescriptorForType();
+    Descriptor descriptor = options.getPrototype().getDescriptorForType();
     if (alreadyBuiltMarshallers.containsKey(descriptor)) {
       return;
     }
 
     TypeDefinition superType =
         TypeDescription.Generic.Builder.parameterizedType(
-                TypeSpecificMarshaller.class, prototype.getClass())
+                TypeSpecificMarshaller.class, options.getPrototype().getClass())
             .build();
 
     // Use default ConstructorStrategy which will generate a constructor with a Message argument of
@@ -208,7 +188,7 @@ private static <T extends Message> void buildOrFindMarshaller(
 
     List<Message> nestedMessagePrototypes = new ArrayList<>();
     for (FieldDescriptor f : descriptor.getFields()) {
-      ProtoFieldInfo field = new ProtoFieldInfo(f, prototype);
+      ProtoFieldInfo field = new ProtoFieldInfo(f, options.getPrototype());
 
       // Store a pre-encoded version of the field variableName to avoid re-encoding all the time.
       String fieldName = CodeGenUtil.fieldNameForSerializedFieldName(field);
@@ -220,7 +200,8 @@ private static <T extends Message> void buildOrFindMarshaller(
                   Modifier.PRIVATE | Modifier.STATIC | Modifier.FINAL)
               .initializer(
                   new SetSerializedFieldName(
-                      fieldName, preservingProtoFieldNames ? f.getName() : f.getJsonName()));
+                      fieldName,
+                      options.isPreservingProtoFieldNames() ? f.getName() : f.getJsonName()));
       if (field.valueJavaType() != JavaType.MESSAGE) {
         continue;
       }
@@ -244,53 +225,60 @@ private static <T extends Message> void buildOrFindMarshaller(
     }
 
     TypeSpecificMarshaller<?> marshaller;
-    buddy =
-        buddy
-            .defineMethod("doMerge", void.class, Modifier.FINAL | Modifier.PROTECTED)
-            .withParameter(JsonParser.class, "parser")
-            .withParameter(int.class, "currentDepth")
-            .withParameter(Message.Builder.class, "messageBuilder")
-            .throwing(IOException.class)
-            .intercept(new DoParse(prototype, ignoringUnknownFields))
-            .defineMethod("doWrite", void.class, Modifier.FINAL | Modifier.PROTECTED)
-            .withParameter(prototype.getClass(), "message")
-            .withParameter(JsonGenerator.class, "gen")
-            .throwing(IOException.class)
-            .intercept(
-                new DoWrite(
-                    prototype,
-                    includingDefaultValueFields,
-                    fieldsToAlwaysOutput,
-                    printingEnumsAsInts,
-                    sortingMapKeys));
-    try {
-      marshaller =
+    TypeSpecificMarshaller<?> cached = MARSHALLER_CACHE.get(options);
+    if (cached != null) {
+      marshaller = cached;
+    } else {
+      buddy =
           buddy
-              .make()
-              .load(TypeSpecificMarshaller.class.getClassLoader())
-              .getLoaded()
-              .getConstructor(prototype.getClass())
-              .newInstance(prototype);
-    } catch (InstantiationException
-        | NoSuchMethodException
-        | InvocationTargetException
-        | IllegalAccessException e) {
-      throw new IllegalStateException(
-          "Could not generate marshaller, this is generally a bug in this library. "
-              + "Please file a report at https://github.com/curioswitch/curiostack with this stack "
-              + "trace and an example proto to reproduce.",
-          e);
+              .defineMethod("doMerge", void.class, Modifier.FINAL | Modifier.PROTECTED)
+              .withParameter(JsonParser.class, "parser")
+              .withParameter(int.class, "currentDepth")
+              .withParameter(Message.Builder.class, "messageBuilder")
+              .throwing(IOException.class)
+              .intercept(new DoParse(options.getPrototype(), options.isIgnoringUnknownFields()))
+              .defineMethod("doWrite", void.class, Modifier.FINAL | Modifier.PROTECTED)
+              .withParameter(options.getPrototype().getClass(), "message")
+              .withParameter(JsonGenerator.class, "gen")
+              .throwing(IOException.class)
+              .intercept(
+                  new DoWrite(
+                      options.getPrototype(),
+                      options.isIncludingDefaultValueFields(),
+                      options.getFieldsToAlwaysOutput(),
+                      options.isPrintingEnumsAsInts(),
+                      options.isSortingMapKeys()));
+      try {
+        marshaller =
+            buddy
+                .make()
+                .load(TypeSpecificMarshaller.class.getClassLoader())
+                .getLoaded()
+                .getConstructor(options.getPrototype().getClass())
+                .newInstance(options.getPrototype());
+        MARSHALLER_CACHE.put(options, marshaller);
+      } catch (InstantiationException
+          | NoSuchMethodException
+          | InvocationTargetException
+          | IllegalAccessException e) {
+        throw new IllegalStateException(
+            "Could not generate marshaller, this is generally a bug in this library. "
+                + "Please file a report at https://github.com/curioswitch/curiostack with this stack "
+                + "trace and an example proto to reproduce.",
+            e);
+      }
     }
     alreadyBuiltMarshallers.put(descriptor, marshaller);
     for (Message nestedPrototype : nestedMessagePrototypes) {
       buildOrFindMarshaller(
-          nestedPrototype,
-          includingDefaultValueFields,
-          fieldsToAlwaysOutput,
-          preservingProtoFieldNames,
-          ignoringUnknownFields,
-          printingEnumsAsInts,
-          sortingMapKeys,
+          new MarshallerOptions(
+              nestedPrototype,
+              options.isIncludingDefaultValueFields(),
+              options.getFieldsToAlwaysOutput(),
+              options.isPreservingProtoFieldNames(),
+              options.isIgnoringUnknownFields(),
+              options.isPrintingEnumsAsInts(),
+              options.isSortingMapKeys()),
           alreadyBuiltMarshallers);
     }
   }