fix: too much memory consumption in LambdaFunctionBootstrap, close #494…

…, #481.
killme2008 · Oct 12, 2022 · 1adb85a · 1adb85a
1 parent a248d80
commit 1adb85a
Show file tree

Hide file tree

Showing 6 changed files with 67 additions and 33 deletions.
diff --git a/src/main/java/com/googlecode/aviator/runtime/LambdaFunctionBootstrap.java b/src/main/java/com/googlecode/aviator/runtime/LambdaFunctionBootstrap.java
@@ -1,5 +1,8 @@
 package com.googlecode.aviator.runtime;
 
+import java.lang.ref.Reference;
+import java.lang.ref.SoftReference;
+import java.lang.ref.WeakReference;
 import java.util.Collection;
 import java.util.Iterator;
 import java.util.List;
@@ -27,7 +30,7 @@ public class LambdaFunctionBootstrap implements Comparable<LambdaFunctionBootstr
   private final List<FunctionParam> params;
   private final boolean inheritEnv;
 
-  private final ThreadLocal<LambdaFunction> fnLocal = new ThreadLocal<>();
+  private final ThreadLocal<Reference<LambdaFunction>> fnLocal = new ThreadLocal<>();
 
 
   @Override
@@ -82,26 +85,22 @@ public Expression getExpression() {
    * @return
    */
   public LambdaFunction newInstance(final Env env) {
-    LambdaFunction fn = null;
-    if (this.inheritEnv && (fn = this.fnLocal.get()) != null) {
-      fn.setContext(env);
-      return fn;
+    Reference<LambdaFunction> ref = null;
+    if (this.inheritEnv && (ref = this.fnLocal.get()) != null) {
+      LambdaFunction fn = ref.get();
+      if (fn != null) {
+        fn.setContext(env);
+        return fn;
+      } else {
+        this.fnLocal.remove();
+      }
     }
 
-    // try {
-    fn = new LambdaFunction(this.name, this.params, this.expression, env);
+    LambdaFunction fn = new LambdaFunction(this.name, this.params, this.expression, env);
     fn.setInheritEnv(this.inheritEnv);
     if (this.inheritEnv) {
-      this.fnLocal.set(fn);
+      this.fnLocal.set(new SoftReference<>(fn));
     }
     return fn;
-    // final LambdaFunction fn =
-    // (LambdaFunction) this.constructor.invoke(this.params, this.expression, env);
-
-    // } catch (ExpressionRuntimeException e) {
-    // throw e;
-    // } catch (Throwable t) {
-    // throw new ExpressionRuntimeException("Fail to create lambda instance.", t);
-    // }
   }
 }
diff --git a/src/main/java/com/googlecode/aviator/runtime/RuntimeUtils.java b/src/main/java/com/googlecode/aviator/runtime/RuntimeUtils.java
@@ -6,6 +6,7 @@
 import com.googlecode.aviator.AviatorEvaluator;
 import com.googlecode.aviator.AviatorEvaluatorInstance;
 import com.googlecode.aviator.Options;
+import com.googlecode.aviator.runtime.function.LambdaFunction;
 import com.googlecode.aviator.runtime.function.internal.UnpackingArgsFunction;
 import com.googlecode.aviator.runtime.type.AviatorFunction;
 import com.googlecode.aviator.runtime.type.AviatorNil;
@@ -57,6 +58,12 @@ public static final AviatorFunction unpackArgsFunction(final AviatorFunction fn)
     return new UnpackingArgsFunction(fn);
   }
 
+  public static void resetLambdaContext(AviatorFunction fn) {
+    if (fn != null && fn instanceof LambdaFunction) {
+      ((LambdaFunction) fn).resetContext();
+    }
+  }
+
   /**
    * Cast an object into sequence if possible, otherwise throw an exception.
    *

diff --git a/src/main/java/com/googlecode/aviator/runtime/function/LambdaFunction.java b/src/main/java/com/googlecode/aviator/runtime/function/LambdaFunction.java
@@ -62,6 +62,14 @@ public void setContext(final Env context) {
     this.context = context;
   }
 
+  public void resetContext() {
+    if (this.inheritEnv) {
+      // gc friendly
+      this.context = null;
+      System.out.println("gogogo");
+    }
+  }
+
   public LambdaFunction(final String name, final List<FunctionParam> params,
       final Expression expression, final Env context) {
     super();

diff --git a/src/main/java/com/googlecode/aviator/runtime/function/internal/IfCallccFunction.java b/src/main/java/com/googlecode/aviator/runtime/function/internal/IfCallccFunction.java
@@ -1,6 +1,7 @@
 package com.googlecode.aviator.runtime.function.internal;
 
 import java.util.Map;
+import com.googlecode.aviator.runtime.RuntimeUtils;
 import com.googlecode.aviator.runtime.function.AbstractFunction;
 import com.googlecode.aviator.runtime.type.AviatorFunction;
 import com.googlecode.aviator.runtime.type.AviatorObject;
@@ -14,7 +15,6 @@
  */
 public class IfCallccFunction extends AbstractFunction {
 
-
   private static final long serialVersionUID = 3511688119189694245L;
 
   private IfCallccFunction() {
@@ -40,12 +40,16 @@ public AviatorObject call(final Map<String, Object> env, final AviatorObject arg
       }
 
       AviatorFunction otherClausesFn = (AviatorFunction) nextClauseVal;
-      AviatorObject result = otherClausesFn.call(env);
-      // No remaining statements, return the if statement result.
-      if (result == Constants.REDUCER_EMPTY) {
-        return arg1;
+      try {
+        AviatorObject result = otherClausesFn.call(env);
+        // No remaining statements, return the if statement result.
+        if (result == Constants.REDUCER_EMPTY) {
+          return arg1;
+        }
+        return result;
+      } finally {
+        RuntimeUtils.resetLambdaContext(otherClausesFn);
       }
-      return result;
     }
   }
 }
diff --git a/src/main/java/com/googlecode/aviator/runtime/function/internal/ReducerFunction.java b/src/main/java/com/googlecode/aviator/runtime/function/internal/ReducerFunction.java
@@ -22,7 +22,6 @@
  */
 public class ReducerFunction extends AbstractFunction {
 
-
   private static final long serialVersionUID = -6117602709327741955L;
 
   private ReducerFunction() {}
@@ -41,6 +40,15 @@ public final AviatorObject call(final Map<String, Object> env, final AviatorObje
     Object coll = arg1.getValue(env);
     AviatorFunction iteratorFn = (AviatorFunction) arg2;
 
+    try {
+      return reduce(env, arg2, arg3, coll, iteratorFn);
+    } finally {
+      RuntimeUtils.resetLambdaContext(iteratorFn);
+    }
+  }
+
+  private AviatorObject reduce(final Map<String, Object> env, final AviatorObject arg2,
+      final AviatorObject arg3, Object coll, AviatorFunction iteratorFn) {
     int maxLoopCount = RuntimeUtils.getInstance(env).getOptionValue(Options.MAX_LOOP_COUNT).number;
     AviatorObject result = AviatorNil.NIL;
     long c = 0;
@@ -134,5 +142,4 @@ public final AviatorObject call(final Map<String, Object> env, final AviatorObje
       return contResult;
     }
   }
-
 }
diff --git a/src/main/java/com/googlecode/aviator/runtime/function/internal/TryCatchFunction.java b/src/main/java/com/googlecode/aviator/runtime/function/internal/TryCatchFunction.java
@@ -1,8 +1,8 @@
 package com.googlecode.aviator.runtime.function.internal;
 
-
 import java.util.List;
 import java.util.Map;
+import com.googlecode.aviator.runtime.RuntimeUtils;
 import com.googlecode.aviator.runtime.function.AbstractFunction;
 import com.googlecode.aviator.runtime.type.AviatorFunction;
 import com.googlecode.aviator.runtime.type.AviatorNil;
@@ -56,9 +56,14 @@ public AviatorObject call(final Map<String, Object> env, final AviatorObject arg
         throw Reflector.sneakyThrow(t);
       }
     } finally {
-      if (finallyBody != null) {
-        AviatorObject ret = finallyBody.call(env);
-        result = chooseResult(result, ret);
+      try {
+        if (finallyBody != null) {
+          AviatorObject ret = finallyBody.call(env);
+          result = chooseResult(result, ret);
+        }
+      } finally {
+        RuntimeUtils.resetLambdaContext(tryBody);
+        RuntimeUtils.resetLambdaContext(finallyBody);
       }
     }
 
@@ -71,11 +76,15 @@ public AviatorObject call(final Map<String, Object> env, final AviatorObject arg
       return result;
     }
     AviatorFunction continueFn = (AviatorFunction) val;
-    AviatorObject contResult = continueFn.call(env);
-    if (contResult == Constants.REDUCER_EMPTY) {
-      return result;
-    } else {
-      return contResult;
+    try {
+      AviatorObject contResult = continueFn.call(env);
+      if (contResult == Constants.REDUCER_EMPTY) {
+        return result;
+      } else {
+        return contResult;
+      }
+    } finally {
+      RuntimeUtils.resetLambdaContext(continueFn);
     }
   }