fix: Close async oneshot jobs on finder:close (#1140)

2 files changed, 23 insertions, 5 deletions

diff --git a/lua/telescope/finders/async_oneshot_finder.lua b/lua/telescope/finders/async_oneshot_finder.lua
index 19c2195..37101fa 100644
--- a/lua/telescope/finders/async_oneshot_finder.lua
+++ b/lua/telescope/finders/async_oneshot_finder.lua
@@ -4,6 +4,8 @@ local LinesPipe = require("telescope._").LinesPipe
 
 local make_entry = require "telescope.make_entry"
 
+local await_count = 1000
+
 return function(opts)
   opts = opts or {}
 
@@ -18,9 +20,14 @@ return function(opts)
   local job_completed = false
   local stdout = nil
 
+  local job
+
   return setmetatable({
-    -- close = function() results = {}; job_started = false end,
-    close = function() end,
+    close = function()
+      if job then
+        job:close()
+      end
+    end,
     results = results,
   }, {
     __call = function(_, prompt, process_result, process_complete)
@@ -36,7 +43,7 @@ return function(opts)
         -- end
 
         stdout = LinesPipe()
-        local _ = async_job.spawn {
+        job = async_job.spawn {
           command = job_opts.command,
           args = job_opts.args,
           cwd = cwd,
@@ -48,9 +55,13 @@ return function(opts)
       end
 
       if not job_completed then
-        for line in stdout:iter(true) do
+        for line in stdout:iter(false) do
           num_results = num_results + 1
 
+          if num_results % await_count then
+            async.util.scheduler()
+          end
+
           local v = entry_maker(line)
           results[num_results] = v
           process_result(v)
@@ -65,7 +76,9 @@ return function(opts)
       local current_count = num_results
       for index = 1, current_count do
         -- TODO: Figure out scheduling...
-        async.util.scheduler()
+        if index % await_count then
+          async.util.scheduler()
+        end
 
         if process_result(results[index]) then
           break
diff --git a/lua/telescope/pickers.lua b/lua/telescope/pickers.lua
index 11fa2fa..7dbe8a0 100644
--- a/lua/telescope/pickers.lua
+++ b/lua/telescope/pickers.lua
@@ -385,6 +385,7 @@ function Picker:find()
     while true do
       -- Wait for the next input
       rx.last()
+      await_schedule()
 
       self:_reset_track()
 
@@ -1125,6 +1126,10 @@ function pickers.on_close_prompt(prompt_bufnr)
     picker.previewer:teardown()
   end
 
+  if picker.finder then
+    picker.finder:close()
+  end
+
   picker.close_windows(status)
 end