Add cleanup hook to nrepl exit (#77)

Then, when the nrepl closes, we cleanup connection-related state to avoid broken channel errors

Author: hkupty

lua/acid/connections.lua

@@ -24,24 +24,33 @@ connections.add = function(addr)
   return ulid
 end
 
-connections.remove = function(addr)
-  local ix
-
-  -- Remove current connections first
-  -- Then remove remaining
-  for pwd, v in pairs(connections.current) do
-    if v[2] == addr[2] and v[1] == addr[1] then
-      ix = connections.current[pwd]
-      connections.store[ix] = nil
-      connections.current[pwd] = nil
+connections.remove = function(pwd, addr)
+  -- If removed address is current to a pwd
+  local key = connections.current[pwd]
+  local conn = connections.store[key]
+
+  -- Double-check if address is correct
+  if key ~= nil and conn ~= nil and conn[2] == addr[2] and conn[1] == addr[1] then
+    -- Then remove it from current pwd
+    connections.current[pwd] = nil
+    -- Remove it's definition
+    connections.store[key] = nil
+
+    -- And remove all other addresses that point to it
+    for ix, v in pairs(connections.current) do
+      if v == key then
+        connections.current[ix] = nil
+      end
     end
-  end
-
-  for i, v in ipairs(connections.store) do
-    if v[2] == addr[2] and v[1] == addr[1] then
-      connections.store[i] = nil
+  else
+    -- Else, remove it from the connections if no address points to it
+    for i, v in pairs(connections.store) do
+      if v[2] == addr[2] and v[1] == addr[1] then
+        connections.store[i] = nil
+      end
     end
   end
+
 end
 
 --- Elects selected connection as primary (thus default) for a certain address

lua/acid/nrepl.lua

@@ -130,6 +130,7 @@ nrepl.start = function(obj)
       cmd , {
         on_stdout = "AcidJobHandler",
         on_stderr = "AcidJobHandler",
+        on_exit = "AcidJobCleanup",
         cwd = pwd
       }
     })
@@ -165,10 +166,21 @@ nrepl.stop = function(obj)
 
   nvim.nvim_call_function("jobstop", {nrepl.cache[pwd].job})
   connections.unselect(pwd)
-  connections.remove(nrepl.cache[pwd].addr)
+  connections.remove(pwd, nrepl.cache[pwd].addr)
   nrepl.cache[obj.pwd] = nil
 end
 
+nrepl.cleanup = function(job_id)
+  nrepl.handle._store[job_id] = nil
+  local pwd = utils.search(nrepl.cache, function(obj)
+    return obj.job == job_id
+  end)
+
+  if pwd ~= nil then
+    nrepl.stop{pwd = pwd}
+  end
+end
+
 nrepl.handle = {
   _store = {},
   stdout = function(dt, ch)

lua/acid/utils.lua

@@ -51,6 +51,22 @@ utils.find = function(coll, val)
   return false
 end
 
+utils.search = function(coll, pred)
+  for key, v in pairs(coll) do
+    if pred(v, key) then
+      return key
+    end
+  end
+  for ix, v in ipairs(coll) do
+    if pred(v, ix) then
+      return ix
+    end
+  end
+  return nil
+end
+
+
+
 utils.map = function(tbl, fn)
   local new = {}
 

plugin/acid.vim

@@ -87,6 +87,11 @@ function! AcidJobHandler(id, data, stream)
   call luaeval('require("acid.nrepl").handle[_A[1]](_A[2], _A[3])', [a:stream, a:data, a:id])
 endfunction
 
+function! AcidJobCleanup(id, data, _)
+  call luaeval('require("acid.nrepl").cleanup(_A[1])', [a:id])
+endfunction
+
+
 map <Plug>(acid-interrupt)        <Cmd>lua require("acid.features").interrupt()<CR>
 map <Plug>(acid-go-to)            <Cmd>lua require("acid.features").go_to()<CR>
 map <Plug>(acid-go-to)            <Cmd>lua require("acid.features").go_to()<CR>