Painless: Add a simple cache for whitelist methods and fields. (#28142)

With support for multiple contexts we are adding some caching to the whitelist to keep the memory footprint for definitions from exploding.

Author: jdconrad

modules/lang-painless/src/main/java/org/elasticsearch/painless/Definition.java

@@ -41,6 +41,9 @@
  */
 public final class Definition {
 
+    private static final Map<String, Method> methodCache = new HashMap<>();
+    private static final Map<String, Field> fieldCache = new HashMap<>();
+
     private static final Pattern TYPE_NAME_PATTERN = Pattern.compile("^[_a-zA-Z][._a-zA-Z0-9]*$");
 
     public static final String[] DEFINITION_FILES = new String[] {
@@ -533,6 +536,22 @@ Collection<Type> allSimpleTypes() {
         return simpleTypesMap.values();
     }
 
+    private static String buildMethodCacheKey(String structName, String methodName, List<Type> arguments) {
+        StringBuilder key = new StringBuilder();
+        key.append(structName);
+        key.append(methodName);
+
+        for (Type argument : arguments) {
+            key.append(argument.name);
+        }
+
+        return key.toString();
+    }
+
+    private static String buildFieldCacheKey(String structName, String fieldName, String typeName) {
+        return structName + fieldName + typeName;
+    }
+
     // INTERNAL IMPLEMENTATION:
 
     private final Map<Class<?>, RuntimeClass> runtimeMap;
@@ -836,8 +855,10 @@ private void addConstructor(String ownerStructName, Whitelist.Constructor whitel
                         " with constructor parameters " + whitelistConstructor.painlessParameterTypeNames);
             }
 
-            painlessConstructor = new Method("<init>", ownerStruct, null, getTypeInternal("void"), painlessParametersTypes,
-                asmConstructor, javaConstructor.getModifiers(), javaHandle);
+            painlessConstructor = methodCache.computeIfAbsent(buildMethodCacheKey(ownerStruct.name, "<init>", painlessParametersTypes),
+                    key -> new Method("<init>", ownerStruct, null, getTypeInternal("void"), painlessParametersTypes,
+                            asmConstructor, javaConstructor.getModifiers(), javaHandle));
+
             ownerStruct.constructors.put(painlessMethodKey, painlessConstructor);
         } else if (painlessConstructor.arguments.equals(painlessParametersTypes) == false){
             throw new IllegalArgumentException(
@@ -859,7 +880,7 @@ private void addMethod(ClassLoader whitelistClassLoader, String ownerStructName,
                     " [" + whitelistMethod.javaMethodName + "] for owner struct [" + ownerStructName + "].");
         }
 
-        Class<?> javaAugmentedClass = null;
+        Class<?> javaAugmentedClass;
 
         if (whitelistMethod.javaAugmentedClassName != null) {
             try {
@@ -869,6 +890,8 @@ private void addMethod(ClassLoader whitelistClassLoader, String ownerStructName,
                         "not found for method with name [" + whitelistMethod.javaMethodName + "] " +
                         "and parameters " + whitelistMethod.painlessParameterTypeNames, cnfe);
             }
+        } else {
+            javaAugmentedClass = null;
         }
 
         int augmentedOffset = javaAugmentedClass == null ? 0 : 1;
@@ -939,8 +962,10 @@ private void addMethod(ClassLoader whitelistClassLoader, String ownerStructName,
                         "[" + whitelistMethod.javaMethodName + "] and parameters " + whitelistMethod.painlessParameterTypeNames);
                 }
 
-                painlessMethod = new Method(whitelistMethod.javaMethodName, ownerStruct, null, painlessReturnType,
-                    painlessParametersTypes, asmMethod, javaMethod.getModifiers(), javaMethodHandle);
+                painlessMethod = methodCache.computeIfAbsent(
+                        buildMethodCacheKey(ownerStruct.name, whitelistMethod.javaMethodName, painlessParametersTypes),
+                        key -> new Method(whitelistMethod.javaMethodName, ownerStruct, null, painlessReturnType, painlessParametersTypes,
+                                asmMethod, javaMethod.getModifiers(), javaMethodHandle));
                 ownerStruct.staticMethods.put(painlessMethodKey, painlessMethod);
             } else if ((painlessMethod.name.equals(whitelistMethod.javaMethodName) && painlessMethod.rtn.equals(painlessReturnType) &&
                     painlessMethod.arguments.equals(painlessParametersTypes)) == false) {
@@ -963,8 +988,10 @@ private void addMethod(ClassLoader whitelistClassLoader, String ownerStructName,
                         "[" + whitelistMethod.javaMethodName + "] and parameters " + whitelistMethod.painlessParameterTypeNames);
                 }
 
-                painlessMethod = new Method(whitelistMethod.javaMethodName, ownerStruct, javaAugmentedClass, painlessReturnType,
-                    painlessParametersTypes, asmMethod, javaMethod.getModifiers(), javaMethodHandle);
+                painlessMethod = methodCache.computeIfAbsent(
+                        buildMethodCacheKey(ownerStruct.name, whitelistMethod.javaMethodName, painlessParametersTypes),
+                        key -> new Method(whitelistMethod.javaMethodName, ownerStruct, javaAugmentedClass, painlessReturnType,
+                                painlessParametersTypes, asmMethod, javaMethod.getModifiers(), javaMethodHandle));
                 ownerStruct.methods.put(painlessMethodKey, painlessMethod);
             } else if ((painlessMethod.name.equals(whitelistMethod.javaMethodName) && painlessMethod.rtn.equals(painlessReturnType) &&
                 painlessMethod.arguments.equals(painlessParametersTypes)) == false) {
@@ -1016,33 +1043,40 @@ private void addField(String ownerStructName, Whitelist.Field whitelistField) {
             Field painlessField = ownerStruct.staticMembers.get(whitelistField.javaFieldName);
 
             if (painlessField == null) {
-                painlessField = new Field(whitelistField.javaFieldName, javaField.getName(),
-                    ownerStruct, painlessFieldType, javaField.getModifiers(), null, null);
+                painlessField = fieldCache.computeIfAbsent(
+                        buildFieldCacheKey(ownerStruct.name, whitelistField.javaFieldName, painlessFieldType.name),
+                        key -> new Field(whitelistField.javaFieldName, javaField.getName(),
+                                ownerStruct, painlessFieldType, javaField.getModifiers(), null, null));
                 ownerStruct.staticMembers.put(whitelistField.javaFieldName, painlessField);
             } else if (painlessField.type.equals(painlessFieldType) == false) {
                 throw new IllegalArgumentException("illegal duplicate static fields [" + whitelistField.javaFieldName + "] " +
                     "found within the struct [" + ownerStruct.name + "] with type [" + whitelistField.painlessFieldTypeName + "]");
             }
         } else {
-            MethodHandle javaMethodHandleGetter = null;
-            MethodHandle javaMethodHandleSetter = null;
+            MethodHandle javaMethodHandleGetter;
+            MethodHandle javaMethodHandleSetter;
 
             try {
                 if (Modifier.isStatic(javaField.getModifiers()) == false) {
                     javaMethodHandleGetter = MethodHandles.publicLookup().unreflectGetter(javaField);
                     javaMethodHandleSetter = MethodHandles.publicLookup().unreflectSetter(javaField);
+                } else {
+                    javaMethodHandleGetter = null;
+                    javaMethodHandleSetter = null;
                 }
             } catch (IllegalAccessException exception) {
                 throw new IllegalArgumentException("getter/setter [" + whitelistField.javaFieldName + "]" +
                     " not found for class [" + ownerStruct.clazz.getName() + "].");
             }
 
-            Field painlessField = ownerStruct.staticMembers.get(whitelistField.javaFieldName);
+            Field painlessField = ownerStruct.members.get(whitelistField.javaFieldName);
 
             if (painlessField == null) {
-                painlessField = new Field(whitelistField.javaFieldName, javaField.getName(),
-                    ownerStruct, painlessFieldType, javaField.getModifiers(), javaMethodHandleGetter, javaMethodHandleSetter);
-                ownerStruct.staticMembers.put(whitelistField.javaFieldName, painlessField);
+                painlessField = fieldCache.computeIfAbsent(
+                        buildFieldCacheKey(ownerStruct.name, whitelistField.javaFieldName, painlessFieldType.name),
+                        key -> new Field(whitelistField.javaFieldName, javaField.getName(),
+                                ownerStruct, painlessFieldType, javaField.getModifiers(), javaMethodHandleGetter, javaMethodHandleSetter));
+                ownerStruct.members.put(whitelistField.javaFieldName, painlessField);
             } else if (painlessField.type.equals(painlessFieldType) == false) {
                 throw new IllegalArgumentException("illegal duplicate member fields [" + whitelistField.javaFieldName + "] " +
                     "found within the struct [" + ownerStruct.name + "] with type [" + whitelistField.painlessFieldTypeName + "]");