Tweak DFA cache cleanup

daniellansun · Nov 27, 2023 · cd67bb2 · cd67bb2
1 parent 697db94
commit cd67bb2
Show file tree

Hide file tree

Showing 3 changed files with 76 additions and 10 deletions.
diff --git a/src/main/java/org/apache/groovy/parser/antlr4/internal/atnmanager/AtnManager.java b/src/main/java/org/apache/groovy/parser/antlr4/internal/atnmanager/AtnManager.java
@@ -20,9 +20,15 @@
 
 import org.antlr.v4.runtime.atn.ATN;
 import org.apache.groovy.util.SystemUtil;
+import org.codehaus.groovy.runtime.DefaultGroovyMethods;
 
+import java.lang.invoke.MethodHandles;
+import java.lang.ref.Reference;
+import java.lang.ref.ReferenceQueue;
+import java.lang.ref.SoftReference;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.logging.Logger;
 
 /**
  * Manage ATN to avoid memory leak
@@ -31,9 +37,18 @@ public abstract class AtnManager {
     private static final ReentrantReadWriteLock RRWL = new ReentrantReadWriteLock(true);
     private static final ReentrantReadWriteLock.WriteLock WRITE_LOCK = RRWL.writeLock();
     public static final ReentrantReadWriteLock.ReadLock READ_LOCK = RRWL.readLock();
+
+    @Deprecated
     private static final String DFA_CACHE_THRESHOLD_OPT = "groovy.antlr4.cache.threshold";
+
+    @Deprecated
     private static final long DFA_CACHE_THRESHOLD;
 
+    private final Logger LOGGER = Logger.getLogger(MethodHandles.lookup().lookupClass().getName());
+
+    private final ReferenceQueue<AtnWrapper> atnWrapperReferenceQueue = new ReferenceQueue<>();
+    private AtnWrapperSoftReference atnWrapperSoftReference;
+
     static {
         long t = SystemUtil.getLongSafe(DFA_CACHE_THRESHOLD_OPT, 64L);
         if (t <= 0) {
@@ -43,7 +58,40 @@ public abstract class AtnManager {
         DFA_CACHE_THRESHOLD = t;
     }
 
-    public abstract ATN getATN();
+    {
+        new Thread(() -> {
+//            LOGGER.info("DFA cleanup thread started");
+            while (true) {
+                try {
+                    Reference<? extends AtnWrapper> reference = atnWrapperReferenceQueue.remove();
+                    if (reference instanceof AtnWrapperSoftReference && shouldClearDfaCache()) {
+                        AtnWrapperSoftReference atnWrapperSoftReference = (AtnWrapperSoftReference) reference;
+                        atnWrapperSoftReference.getAtnWrapper().clearDFA();
+                        LOGGER.info("DFA cache cleared");
+                    }
+                } catch (InterruptedException e) {
+                    LOGGER.warning(DefaultGroovyMethods.asString(e));
+                }
+            }
+        }).start();
+    }
+
+    public ATN getATN() {
+        return getAtnWrapper().getATN();
+    }
+
+    protected abstract AtnWrapper createAtnWrapper();
+
+    protected synchronized AtnWrapper getAtnWrapper() {
+        AtnWrapper atnWrapper;
+
+        if (null == atnWrapperSoftReference || null == (atnWrapper = atnWrapperSoftReference.get())) {
+            atnWrapper = createAtnWrapper();
+            atnWrapperSoftReference = new AtnWrapperSoftReference(atnWrapper, atnWrapperReferenceQueue);
+        }
+
+        return atnWrapper;
+    }
 
     protected abstract boolean shouldClearDfaCache();
 
@@ -55,6 +103,7 @@ public AtnWrapper(ATN atn) {
             this.atn = atn;
         }
 
+        @Deprecated
         public ATN checkAndClear() {
             if (!shouldClearDfaCache()) {
                 return atn;
@@ -64,14 +113,35 @@ public ATN checkAndClear() {
                 return atn;
             }
 
+            clearDFA();
+
+            return atn;
+        }
+
+        public ATN getATN() {
+            return atn;
+        }
+
+        public void clearDFA() {
             WRITE_LOCK.lock();
             try {
                 atn.clearDFA();
             } finally {
                 WRITE_LOCK.unlock();
             }
+        }
+    }
 
-            return atn;
+    protected static class AtnWrapperSoftReference extends SoftReference<AtnWrapper> {
+        private AtnWrapper atnWrapper;
+
+        public AtnWrapperSoftReference(AtnWrapper referent, ReferenceQueue<? super AtnWrapper> q) {
+            super(referent, q);
+            this.atnWrapper = atnWrapper;
+        }
+
+        public AtnWrapper getAtnWrapper() {
+            return atnWrapper;
         }
     }
 }
diff --git a/src/main/java/org/apache/groovy/parser/antlr4/internal/atnmanager/LexerAtnManager.java b/src/main/java/org/apache/groovy/parser/antlr4/internal/atnmanager/LexerAtnManager.java
@@ -18,7 +18,6 @@
  */
 package org.apache.groovy.parser.antlr4.internal.atnmanager;
 
-import org.antlr.v4.runtime.atn.ATN;
 import org.apache.groovy.parser.antlr4.GroovyLangLexer;
 import org.apache.groovy.util.SystemUtil;
 
@@ -28,16 +27,15 @@
 public class LexerAtnManager extends AtnManager {
     private static final String GROOVY_CLEAR_LEXER_DFA_CACHE = "groovy.antlr4.clear.lexer.dfa.cache";
     private static final boolean TO_CLEAR_LEXER_DFA_CACHE;
-    private final AtnWrapper lexerAtnWrapper = new AtnManager.AtnWrapper(GroovyLangLexer._ATN);
     public static final LexerAtnManager INSTANCE = new LexerAtnManager();
 
     static {
         TO_CLEAR_LEXER_DFA_CACHE = SystemUtil.getBooleanSafe(GROOVY_CLEAR_LEXER_DFA_CACHE);
     }
 
     @Override
-    public ATN getATN() {
-        return lexerAtnWrapper.checkAndClear();
+    protected AtnWrapper createAtnWrapper() {
+        return new AtnManager.AtnWrapper(GroovyLangLexer._ATN);
     }
 
     @Override

diff --git a/src/main/java/org/apache/groovy/parser/antlr4/internal/atnmanager/ParserAtnManager.java b/src/main/java/org/apache/groovy/parser/antlr4/internal/atnmanager/ParserAtnManager.java
@@ -18,19 +18,17 @@
  */
 package org.apache.groovy.parser.antlr4.internal.atnmanager;
 
-import org.antlr.v4.runtime.atn.ATN;
 import org.apache.groovy.parser.antlr4.GroovyLangParser;
 
 /**
  * Manage ATN for parser to avoid memory leak
  */
 public class ParserAtnManager extends AtnManager {
-    private final AtnWrapper parserAtnWrapper = new AtnManager.AtnWrapper(GroovyLangParser._ATN);
     public static final ParserAtnManager INSTANCE = new ParserAtnManager();
 
     @Override
-    public ATN getATN() {
-        return parserAtnWrapper.checkAndClear();
+    protected AtnWrapper createAtnWrapper() {
+        return new AtnManager.AtnWrapper(GroovyLangParser._ATN);
     }
 
     @Override