fix: resolve benchmark cold start OOM by resetting DFA cache

benchmark/sqlBenchmark.ts

@@ -110,35 +110,30 @@ class SqlBenchmark {
 
     private _lastResultsCache: BenchmarkResult[] | null = null;
 
+    private _parserConstructor: (new () => BasicSQL) | null = null;
+
     /**
      * Returns SqlParser instance of specific language.
      */
     getSqlParser(): BasicSQL {
+        if (!this._parserConstructor) {
+            this._parserConstructor = require(path.resolve(`src/parser/${this.language}/index.ts`))[
+                languageNameMap[this.language]
+            ];
+        }
+        const parser = new this._parserConstructor();
         if (!this.isHot) {
-            this.clearATNCache();
+            parser.resetDfaCache();
         }
-        const Parser = require(path.resolve(`src/parser/${this.language}/index.ts`))[
-            languageNameMap[this.language]
-        ];
-        return new Parser();
+        return parser;
     }
 
     /**
-     * Due to the presence of ATN cache in antlr4, we will clear the module cache to ensure that each parser is brand new and with no cache.
+     * Reset the shared DFA caches before collecting hot-start samples.
      */
-    clearATNCache() {
-        const caches = Object.keys(require.cache);
-        const sourcePath = path.join(__dirname, '../src');
-        caches
-            .filter((cache) => cache.includes(sourcePath))
-            .forEach((moduleName) => {
-                const module = require.cache[moduleName]!;
-                // Fix Memory Leak
-                if (module.parent) {
-                    module.parent.children = [];
-                }
-                delete require.cache[moduleName];
-            });
+    resetDfaCache() {
+        const parser = this.getSqlParser();
+        parser.resetDfaCache();
     }
 
     getColumns = () => {
@@ -166,7 +161,7 @@ class SqlBenchmark {
             this._lastResultsCache?.find((item) => item.type === type && item.name === name) ?? {};
 
         if (this.isHot) {
-            this.clearATNCache();
+            this.resetDfaCache();
         }
 
         if (this.isHot && loopTimes < 2) {

package.json

@@ -29,14 +29,14 @@
     "build": "rm -rf dist && tsc",
     "watch": "tsc -w",
     "check-types": "tsc -p ./tsconfig.json && tsc -p ./test/tsconfig.json",
-    "test": "NODE_OPTIONS=--max_old_space_size=4096 && jest",
+    "test": "NODE_OPTIONS=--max_old_space_size=4096 jest",
     "release": "node ./scripts/release.js",
     "prettier-check": "prettier --check .",
     "format": "prettier --write .",
     "format-g4": "antlr-format -c ./antlr.format.json -v ./src/grammar/**/*.g4",
     "cleanComment": "node ./scripts/cleanCommentCli.js",
-    "benchmark": "NODE_OPTIONS=--max_old_space_size=4096 && node --no-warnings ./scripts/benchmark.js",
-    "benchmark:release": "node --no-warnings ./scripts/benchmark.js --release"
+    "benchmark": "NODE_OPTIONS=--max_old_space_size=4096 node --no-warnings ./scripts/benchmark.js",
+    "benchmark:release": "NODE_OPTIONS=--max_old_space_size=4096 node --no-warnings ./scripts/benchmark.js --release"
   },
   "license": "MIT",
   "dependencies": {

scripts/benchmark.js

@@ -1,3 +1,4 @@
+process.env.NODE_OPTIONS = process.env.NODE_OPTIONS || '--max_old_space_size=4096';
 const path = require('path');
 const spawn = require('child_process').spawn;
 const fs = require('fs');
@@ -108,7 +109,7 @@ function main() {
     if (argv.lang) {
         const supportedLanguage = languages.some((language) => language === argv.lang);
         if (supportedLanguage || argv.lang === 'All Languages') {
-            runBenchmark(argv.lang);
+            runBenchmark(argv.lang, argv.hot ? 'hot' : 'cold');
         } else {
             console.error(
                 chalk.bold.red('\n[Invalid language]:'),

src/parser/common/basicSQL.ts

@@ -270,6 +270,22 @@ export abstract class BasicSQL<
         return this._parseTree;
     }
 
+    /**
+     * Reset the shared Lexer and Parser DFA caches for cold-start benchmarks.
+     * This affects every parser instance of the same dialect.
+     */
+    public resetDfaCache(): void {
+        if (!this._parser || !this._lexer) {
+            this.createParserWithCache('');
+        }
+        this._parser.interpreter.clearDFA();
+        this._lexer.interpreter.clearDFA();
+        this._parseTree = null;
+        this._parsedInput = '';
+        this._parseErrors = [];
+        this._statementStartTokenTypes = null;
+    }
+
     /**
      * Validate input string and return syntax errors if exists.
      * When input contains multiple statements separated by semicolons

test/common/basicSQL.test.ts

@@ -3,6 +3,11 @@ import { FlinkSqlLexer } from 'src/lib/flink/FlinkSqlLexer';
 import { ErrorListener } from 'src/parser/common/parseErrorListener';
 import { CommonTokenStream } from 'antlr4ng';
 
+type DfaInterpreter = { decisionToDFA: Array<{ states: { length: number } }> };
+
+const countDfaStates = (interpreter: DfaInterpreter) =>
+    interpreter.decisionToDFA.reduce((count, dfa) => count + dfa.states.length, 0);
+
 describe('BasicSQL unit tests', () => {
     const flink = new FlinkSQL();
     test('Create lexer', () => {
@@ -81,6 +86,31 @@ describe('BasicSQL unit tests', () => {
         expect(errors.length).not.toBe(0);
     });
 
+    test('Reset shared DFA caches and reparse the same input', () => {
+        const parser = new FlinkSQL();
+        const sql = 'SELECT a, b FROM tbl1 WHERE a = 1 ORDER BY b';
+
+        expect(parser.validate(sql)).toHaveLength(0);
+
+        const parserInternals = parser as unknown as {
+            _lexer: { interpreter: DfaInterpreter };
+            _parser: { interpreter: DfaInterpreter };
+        };
+
+        expect(countDfaStates(parserInternals._lexer.interpreter)).toBeGreaterThan(0);
+        expect(countDfaStates(parserInternals._parser.interpreter)).toBeGreaterThan(0);
+
+        parser.resetDfaCache();
+
+        expect(countDfaStates(parserInternals._lexer.interpreter)).toBe(0);
+        expect(countDfaStates(parserInternals._parser.interpreter)).toBe(0);
+        expect(parser.getParsedInput()).toBe('');
+
+        expect(parser.validate(sql)).toHaveLength(0);
+        expect(countDfaStates(parserInternals._lexer.interpreter)).toBeGreaterThan(0);
+        expect(countDfaStates(parserInternals._parser.interpreter)).toBeGreaterThan(0);
+    });
+
     test('Get All tokens', () => {
         const sql = 'SELECT * FROM tbl1;';
         const tokens = flink.getAllTokens(sql);