fix(treesitter): logger memory leak

4 files changed, 29 insertions, 14 deletions

diff --git a/runtime/lua/vim/treesitter/languagetree.lua b/runtime/lua/vim/treesitter/languagetree.lua
index f274edf961..b4c9027794 100644
--- a/runtime/lua/vim/treesitter/languagetree.lua
+++ b/runtime/lua/vim/treesitter/languagetree.lua
@@ -156,6 +156,7 @@ function LanguageTree:_set_logger()
 
   local lang = self:lang()
 
+  vim.fn.mkdir(vim.fn.stdpath('log'), 'p')
   local logfilename = vim.fs.joinpath(vim.fn.stdpath('log'), 'treesitter.log')
 
   local logfile, openerr = io.open(logfilename, 'a+')
@@ -208,7 +209,8 @@ function LanguageTree:_log(...)
 
   local info = debug.getinfo(2, 'nl')
   local nregions = #self:included_regions()
-  local prefix = string.format('%s:%d: (#regions=%d) ', info.name, info.currentline, nregions)
+  local prefix =
+    string.format('%s:%d: (#regions=%d) ', info.name or '???', info.currentline or 0, nregions)
 
   local msg = { prefix }
   for _, x in ipairs(args) do
diff --git a/src/nvim/lua/executor.c b/src/nvim/lua/executor.c
index 9215926434..459c69f385 100644
--- a/src/nvim/lua/executor.c
+++ b/src/nvim/lua/executor.c
@@ -1302,6 +1302,9 @@ LuaRef nlua_ref(lua_State *lstate, nlua_ref_state_t *ref_state, int index)
   return ref;
 }
 
+// TODO(lewis6991): Currently cannot be run in __gc metamethods as they are
+// invoked in lua_close() which can be invoked after the ref_markers map is
+// destroyed in nlua_common_free_all_mem.
 LuaRef nlua_ref_global(lua_State *lstate, int index)
 {
   return nlua_ref(lstate, nlua_global_refs, index);
diff --git a/src/nvim/lua/treesitter.c b/src/nvim/lua/treesitter.c
index 66a75f8d40..1e559316dd 100644
--- a/src/nvim/lua/treesitter.c
+++ b/src/nvim/lua/treesitter.c
@@ -325,6 +325,17 @@ static TSParser **parser_check(lua_State *L, uint16_t index)
   return luaL_checkudata(L, index, TS_META_PARSER);
 }
 
+static void logger_gc(TSLogger logger)
+{
+  if (!logger.log) {
+    return;
+  }
+
+  TSLuaLoggerOpts *opts = (TSLuaLoggerOpts *)logger.payload;
+  luaL_unref(opts->lstate, LUA_REGISTRYINDEX, opts->cb);
+  xfree(opts);
+}
+
 static int parser_gc(lua_State *L)
 {
   TSParser **p = parser_check(L, 1);
@@ -332,12 +343,7 @@ static int parser_gc(lua_State *L)
     return 0;
   }
 
-  TSLogger logger = ts_parser_logger(*p);
-  if (logger.log) {
-    TSLuaLoggerOpts *opts = (TSLuaLoggerOpts *)logger.payload;
-    xfree(opts);
-  }
-
+  logger_gc(ts_parser_logger(*p));
   ts_parser_delete(*p);
   return 0;
 }
@@ -698,7 +704,7 @@ static void logger_cb(void *payload, TSLogType logtype, const char *s)
 
   lua_State *lstate = opts->lstate;
 
-  nlua_pushref(lstate, opts->cb);
+  lua_rawgeti(lstate, LUA_REGISTRYINDEX, opts->cb);
   lua_pushstring(lstate, logtype == TSLogTypeParse ? "parse" : "lex");
   lua_pushstring(lstate, s);
   if (lua_pcall(lstate, 2, 0, 0)) {
@@ -726,11 +732,13 @@ static int parser_set_logger(lua_State *L)
   }
 
   TSLuaLoggerOpts *opts = xmalloc(sizeof(TSLuaLoggerOpts));
+  lua_pushvalue(L, 4);
+  LuaRef ref = luaL_ref(L, LUA_REGISTRYINDEX);
 
   *opts = (TSLuaLoggerOpts){
     .lex = lua_toboolean(L, 2),
     .parse = lua_toboolean(L, 3),
-    .cb = nlua_ref_global(L, 4),
+    .cb = ref,
     .lstate = L
   };
 
@@ -753,7 +761,7 @@ static int parser_get_logger(lua_State *L)
   TSLogger logger = ts_parser_logger(*p);
   if (logger.log) {
     TSLuaLoggerOpts *opts = (TSLuaLoggerOpts *)logger.payload;
-    nlua_pushref(L, opts->cb);
+    lua_rawgeti(L, LUA_REGISTRYINDEX, opts->cb);
   } else {
     lua_pushnil(L);
   }
diff --git a/test/functional/treesitter/parser_spec.lua b/test/functional/treesitter/parser_spec.lua
index cc833a67cc..ae3b0483c5 100644
--- a/test/functional/treesitter/parser_spec.lua
+++ b/test/functional/treesitter/parser_spec.lua
@@ -9,10 +9,13 @@ local pcall_err = helpers.pcall_err
 local feed = helpers.feed
 local is_os = helpers.is_os
 
-before_each(clear)
-
 describe('treesitter parser API', function()
-  clear()
+  before_each(function()
+    clear()
+    exec_lua[[
+      vim.g.__ts_debug = 1
+    ]]
+  end)
 
   it('parses buffer', function()
     insert([[
@@ -629,7 +632,6 @@ int x = INT_MAX;
     describe("when parsing regions independently", function()
       it("should inject a language", function()
         exec_lua([[
-        vim.g.__ts_debug = 1
         parser = vim.treesitter.get_parser(0, "c", {
           injections = {
             c = "(preproc_def (preproc_arg) @c) (preproc_function_def value: (preproc_arg) @c)"}})