resty/concurrent/timer_task: add possibility of running for the last …

…time

gateway/src/resty/concurrent/timer_task.lua

@@ -9,6 +9,11 @@ local default_interval_seconds = 60
 
 _M.active_tasks = {}
 
+-- Whether a run should be the last one for a given ID
+-- When a task is marked to run for a last time, it will do so even if it is
+-- cancelled.
+_M.last_one = {}
+
 function _M.register_task(id)
   _M.active_tasks[id] = true
 end
@@ -58,13 +63,19 @@ end
 local run_periodic, schedule_next, timer_execute
 
 run_periodic = function(run_now, id, func, args, interval)
-  if not _M.task_is_active(id) then return end
+  if not _M.task_is_active(id) and not _M.last_one[id] then
+    return
+  end
 
   if run_now then
     func(unpack(args))
   end
 
-  schedule_next(id, func, args, interval)
+  if not _M.last_one[id] then
+    schedule_next(id, func, args, interval)
+  else
+    _M.last_one[id] = nil
+  end
 end
 
 -- Note: ngx.timer.at always sends "premature" as the first param.
@@ -89,7 +100,17 @@ function _M:execute(run_now)
   run_periodic(run_now or false, self.id, self.task, self.args, self.interval)
 end
 
-function _M:cancel()
+--- Cancel a task
+-- @tparam[opt] run_one_more boolean True to ensure that the task will run one
+--   more time before it is cancelled. False to just cancel the task. (Defaults
+--   to false)
+function _M:cancel(run_one_more)
+  if run_one_more then
+    _M.last_one[self.id] = true
+  end
+
+  -- We can cancel the task in all cases because the flag to run for the last
+  -- time has precedence.
   _M.unregister_task(self.id)
 end
 

spec/resty/concurrent/timer_task_spec.lua

@@ -6,6 +6,7 @@ describe('TimerTask', function()
 
   before_each(function()
     TimerTask.active_tasks = {}
+    TimerTask.last_one = {}
   end)
 
   after_each(function()
@@ -85,6 +86,14 @@ describe('TimerTask', function()
 
       assert.is_false(TimerTask.task_is_active(task.id))
     end)
+
+    it('marks the task to run for the last time when specified in the params', function()
+      local task = TimerTask.new(test_task)
+
+      task:cancel(true)
+
+      assert.is_true(TimerTask.last_one[task.id])
+    end)
   end)
 
   describe(':execute', function()
@@ -161,6 +170,27 @@ describe('TimerTask', function()
         assert.stub(ngx_timer_stub).was_called()
       end)
     end)
+
+    describe('when the task should run for the last time', function()
+      it('runs the task', function()
+        local timer_task = TimerTask.new(func, { args = args, interval = interval })
+        local func_spy = spy.on(timer_task, 'task')
+        timer_task:cancel(true)
+
+        timer_task:execute(true)
+
+        assert.spy(func_spy).was_called_with(unpack(args))
+      end)
+
+      it('does not schedule another task', function()
+        local timer_task = TimerTask.new(func, { args = args, interval = interval })
+        timer_task:cancel(true)
+
+        timer_task:execute(true)
+
+        assert.stub(ngx_timer_stub).was_not_called()
+      end)
+    end)
   end)
 
   it('cancels itself when it is garbage collected', function()
@@ -172,4 +202,14 @@ describe('TimerTask', function()
 
     assert.is_false(TimerTask.task_is_active(id))
   end)
+
+  it('does not ensure a last run when garbage collected', function()
+    local timer_task = TimerTask.new(test_task)
+    local id = timer_task.id
+
+    timer_task = nil
+    collectgarbage()
+
+    assert.is_falsy(TimerTask.last_one[id])
+  end)
 end)