From 7ea65811caa0b8bea3f80ba017c4b7dc803ba0ed Mon Sep 17 00:00:00 2001 From: Jonas Brusman Date: Tue, 1 Sep 2026 10:15:11 +0200 Subject: [PATCH] Add limit option for bounded prefetch and drain for graceful shutdown --- lib/async/job/processor/redis/server.rb | 93 +++++++++++++++++++------ test/async/job/processor/redis/limit.rb | 84 ++++++++++++++++++++++ 2 files changed, 155 insertions(+), 22 deletions(-) create mode 100644 test/async/job/processor/redis/limit.rb diff --git a/lib/async/job/processor/redis/server.rb b/lib/async/job/processor/redis/server.rb index d3ead22..011d701 100644 --- a/lib/async/job/processor/redis/server.rb +++ b/lib/async/job/processor/redis/server.rb @@ -3,9 +3,11 @@ # Released under the MIT License. # Copyright, 2024-2025, by Samuel Williams. +require "async/barrier" require "async/idler" require "async/job/coder" require "async/job/processor/generic" +require "async/semaphore" require "securerandom" @@ -29,33 +31,46 @@ class Server < Generic # @parameter coder [Async::Job::Coder] The job serialization codec. # @parameter resolution [Integer] The resolution in seconds for delayed job processing. # @parameter parent [Async::Task] The parent task for background processing. - def initialize(delegate, client, prefix: "async-job", coder: Coder::DEFAULT, resolution: 10, parent: nil) + # @parameter limit [Integer | Nil] The maximum number of jobs claimed and processed at once, or nil for no limit. + def initialize(delegate, client, prefix: "async-job", coder: Coder::DEFAULT, resolution: 10, parent: nil, limit: nil) super(delegate) - + @id = SecureRandom.uuid @client = client @prefix = prefix @coder = coder @resolution = resolution - + @job_store = JobStore.new(@client, "#{@prefix}:jobs") @delayed_jobs = DelayedJobs.new(@client, "#{@prefix}:delayed") @ready_list = ReadyList.new(@client, "#{@prefix}:ready") @processing_list = ProcessingList.new(@client, "#{@prefix}:processing", @id, @ready_list, @job_store) - + @parent = parent || Async::Idler.new + + # Limits how many jobs are claimed (moved into this server's pending list) and processed at once. Without a limit, jobs are fetched as fast as they arrive and buffered in this process, which hides the backlog from the ready list, grows memory with the backlog, and abandons every buffered job on shutdown. + @semaphore = limit && Async::Semaphore.new(limit) + + # Jobs run on their own tasks, tracked separately from the fetch loop so that stopping the loop (see #drain) does not cancel jobs that are already running. + @jobs = Async::Barrier.new end + + # @attribute [Async::Semaphore | Nil] The concurrency limit semaphore, if a limit was given. + attr :semaphore # Start the job processing loop immediately. # @returns [Async::Task | false] The processing task or false if already started. def start! return false if @task - + @task = true - + + # Host job tasks outside the fetch loop's task tree, so that stopping the loop (see #drain) does not cancel running jobs: + @job_host ||= @parent.async(transient: true, annotation: "#{self.class.name} jobs") {sleep} + @parent.async(transient: true, annotation: self.class.name) do |task| @task = task - + while true self.dequeue(task) end @@ -78,12 +93,37 @@ def start self.start! end - # Stop the server and all background processing tasks. + # Stop the server and all background processing tasks, including any running jobs. def stop @task&.stop - + @jobs.stop + @job_host&.stop + @job_host = nil + super end + + # Stop fetching new jobs and wait for running jobs to finish. + # + # Jobs that do not finish within the timeout are left running; a subsequent {stop} cancels them and they will be recovered as abandoned jobs. Without a timeout, waits indefinitely. + # + # @parameter timeout [Numeric | Nil] The maximum time to wait for running jobs to finish. + # @returns [Boolean] True if all running jobs finished. + def drain(timeout: nil) + @task&.stop + + if timeout + Task.current.with_timeout(timeout) do + @jobs.wait + true + rescue Async::TimeoutError + false + end + else + @jobs.wait + true + end + end # Generates a human-readable string representing the current statistics. # @@ -117,24 +157,33 @@ def call(job) # Dequeue a job from the ready list and process it. # + # If a limit was given, waits for a slot before claiming the next job, so at most `limit` jobs are claimed at once and any backlog stays in the ready list. + # # If the job fails for any reason, it will be retried. # # If you do not desire this behavior, you should catch exceptions in the delegate. def dequeue(parent) - _id = @processing_list.fetch - - parent.async do - id = _id; _id = nil - - job = @coder.load(@job_store.get(id)) - @delegate.call(job) - @processing_list.complete(id) - rescue => error - Console.error(self, "Job failed with error!", id: id, exception: error) - @processing_list.retry(id) + @semaphore&.acquire + + begin + id = @processing_list.fetch + + @jobs.async(parent: @job_host) do + job = @coder.load(@job_store.get(id)) + @delegate.call(job) + @processing_list.complete(id) + rescue => error + Console.error(self, "Job failed with error!", id: id, exception: error) + @processing_list.retry(id) + ensure + @semaphore&.release + end + rescue Exception + # The job (if claimed) was never handed to a task; put it back and release the slot: + @processing_list.retry(id) if id + @semaphore&.release + raise end - ensure - @processing_list.retry(_id) if _id end private diff --git a/test/async/job/processor/redis/limit.rb b/test/async/job/processor/redis/limit.rb new file mode 100644 index 0000000..a895e97 --- /dev/null +++ b/test/async/job/processor/redis/limit.rb @@ -0,0 +1,84 @@ +# frozen_string_literal: true + +# Released under the MIT License. +# Copyright, 2026, by Samuel Williams. + +require "async" +require "async/redis" + +require "sus/fixtures/async/reactor_context" +require "sus/fixtures/console" + +require "async/job/buffer" +require "async/job/processor/redis" + +describe Async::Job::Processor::Redis do + include Sus::Fixtures::Async::ReactorContext + include Sus::Fixtures::Console::CapturedLogger + + let(:buffer) {Async::Job::Buffer.new} + + let(:prefix) {"async:job:#{SecureRandom.hex(8)}"} + let(:server) {subject.new(buffer, prefix:, resolution: 1, limit: 1)} + + before do + server.start + end + + after do + server.stop + end + + let(:job) {{"data" => "test job"}} + + with "a concurrency limit" do + it "claims at most limit jobs at a time" do + current = 0 + peak = 0 + + mock(buffer) do |mock| + mock.before(:call) do |job| + current += 1 + peak = [peak, current].max + sleep(0.01) + current -= 1 + end + end + + 4.times {|index| server.call(job.merge("index" => index))} + + 4.times {buffer.pop} + + expect(peak).to be == 1 + end + end + + with "#drain" do + it "stops fetching but waits for running jobs to finish" do + started = 0 + + mock(buffer) do |mock| + mock.before(:call) do |job| + started += 1 + sleep(0.05) + end + end + + 4.times {|index| server.call(job.merge("index" => index))} + + # Wait for the first job to start: + until started == 1 + sleep(0.001) + end + + expect(server.drain(timeout: 5)).to be == true + + # The running job finished: + expect(buffer.pop).to have_keys("data" => be == job["data"]) + + # No further jobs were claimed after the drain: + sleep(0.1) + expect(started).to be == 1 + end + end +end