fix(token): ensure that running states are guarded by run tokens

Author: asmodeus812

lua/fuzzy/match.lua

@@ -335,6 +335,8 @@ function Match:match(list, pattern, callback, transform)
     end
 
     -- init core match context
+    self._state.token = (self._state.token or 0) + 1
+    local token = self._state.token
     self.list = assert(list)
     self.pattern = assert(pattern)
     self.callback = assert(callback)
@@ -390,9 +392,12 @@ function Match:match(list, pattern, callback, transform)
     self._state.timer = vim.loop.new_timer()
     self._state.timer:start(0,
         self._options.timer,
-        vim.schedule_wrap(self:_bind_method(
-            Match._match_worker
-        ))
+        vim.schedule_wrap(function()
+            if self._state.token ~= token then
+                return
+            end
+            Match._match_worker(self)
+        end)
     )
 
     -- run one cycle immediately now

lua/fuzzy/stream.lua

@@ -419,6 +419,9 @@ function Stream:start(cmd, opts)
         end
     end
 
+    self._state.token = (self._state.token or 0) + 1
+    local token = self._state.token
+
     self.transform = opts.transform
     self.callback = assert(opts.callback)
 
@@ -441,6 +444,9 @@ function Stream:start(cmd, opts)
     if type(cmd) == "function" then
         local did_finalize = false
         local callback = function(data)
+            if self._state.token ~= token then
+                return
+            end
             if not did_finalize and data ~= nil then
                 self:_handle_data(
                     data,
@@ -460,9 +466,11 @@ function Stream:start(cmd, opts)
             )
             local code = not ok and 1 or 0
             if did_finalize == false then
-                self._state.stdouteof = true
-                self._state.stderreof = true
-                stream:_handle_out(code, err, 3)
+                if stream._state.token == token then
+                    self._state.stdouteof = true
+                    self._state.stderreof = true
+                    stream:_handle_out(code, err, 3)
+                end
                 did_finalize = true
             end
         end)
@@ -479,22 +487,28 @@ function Stream:start(cmd, opts)
             env = opts.env,
             stdio = stdio,
             hide = true,
-        }, self:_bind_method(
-            Stream._handle_exit)
-        ))
+        }, function(...)
+            if self._state.token == token then
+                return Stream._handle_exit(self, ...)
+            end
+        end))
 
         vim.loop.read_start(
             self._state.stdout,
-            self:_bind_method(
-                Stream._handle_stdout
-            )
+            function(...)
+                if self._state.token == token then
+                    return Stream._handle_stdout(self, ...)
+                end
+            end
         )
 
         vim.loop.read_start(
             self._state.stderr,
-            self:_bind_method(
-                Stream._handle_stderr
-            )
+            function(...)
+                if self._state.token == token then
+                    return Stream._handle_stderr(self, ...)
+                end
+            end
         )
     end
 

tests/match_spec.lua

@@ -227,6 +227,35 @@ local function run_wait_timeout_case()
     helpers.assert_ok(saw_results, "wait timeout saw results")
 end
 
+local function run_restart_token_case()
+    local match = new_match({ step = 1, timer = 50 })
+    local list = {}
+    for i = 1, 5000 do
+        list[i] = "item-" .. i
+    end
+    list[100] = "alpha"
+    list[200] = "beta"
+
+    local late_calls = 0
+    local blocked = false
+
+    match:match(list, "alpha", function(results)
+        if blocked and results ~= nil then
+            late_calls = late_calls + 1
+        end
+    end)
+
+    helpers.assert_ok(helpers.wait_for(function()
+        return match:running()
+    end, 500), "restart token running")
+
+    blocked = true
+    match:match(list, "beta", function() end)
+    match:wait(1500)
+
+    helpers.eq(late_calls, 0, "restart token stale callback")
+end
+
 local function run_restart_ephemeral_true_case()
     local match = new_match({ step = 1, timer = 2, ephemeral = true })
     local first = {}
@@ -345,6 +374,7 @@ function M.run()
     helpers.run_test_case("match_multi_chunk", run_multi_chunk_case)
     helpers.run_test_case("match_transform_function", run_transform_function_case)
     helpers.run_test_case("match_wait_timeout", run_wait_timeout_case)
+    helpers.run_test_case("match_restart_token", run_restart_token_case)
     helpers.run_test_case("match_restart_ephemeral_true", run_restart_ephemeral_true_case)
     helpers.run_test_case("match_restart_ephemeral_false", run_restart_ephemeral_false_case)
     helpers.run_test_case("match_tail_chunk", run_tail_chunk_case)

tests/stream_spec.lua

@@ -166,6 +166,50 @@ local function run_restart_case()
     stream:destroy()
 end
 
+local function run_restart_token_case()
+    local Async = require("fuzzy.async")
+    local stream = Stream.new({
+        lines = true,
+        step = 2,
+        ephemeral = true,
+    })
+    local state = { allow_old_finish = false }
+
+    stream:start(function(cb)
+        cb("old-1")
+        while not state.allow_old_finish do
+            Async.yield()
+        end
+        cb("old-2")
+        cb(nil)
+    end, {
+        callback = function() end,
+    })
+
+    helpers.assert_ok(helpers.wait_for(function()
+        return stream:running()
+    end, 500), "restart token running")
+
+    stream:start(function(cb)
+        cb("new-1")
+        cb("new-2")
+        cb(nil)
+    end, {
+        callback = function() end,
+    })
+
+    state.allow_old_finish = true
+    helpers.assert_ok(helpers.wait_for(function()
+        return not stream:running()
+    end, 1500), "restart token done")
+
+    local results = stream.results or {}
+    helpers.eq(#results, 2, "restart token count")
+    helpers.eq(results[1], "new-1", "restart token first")
+    helpers.eq(results[2], "new-2", "restart token second")
+    stream:destroy()
+end
+
 local function run_context_case()
     local stream = Stream.new({ lines = true, step = 1 })
     helpers.assert_ok(vim.fn.executable("awk") == 1, "awk missing")
@@ -480,6 +524,7 @@ function M.run()
     helpers.run_test_case("stream_bytes", run_bytes_case)
     helpers.run_test_case("stream_trim_results_case", run_trim_results_case)
     helpers.run_test_case("stream_restart", run_restart_case)
+    helpers.run_test_case("stream_restart_token", run_restart_token_case)
     helpers.run_test_case("stream_context", run_context_case)
     helpers.run_test_case("stream_partial_line", run_partial_line_case)
     helpers.run_test_case("stream_trailing_partial", run_trailing_partial_case)