Skip to content

Commit 31479ef

Browse files
clauderobacarp
authored andcommitted
Add ConcurrencyLimitedDequeueAdapter for per-queue executor scaling
Introduces a new dequeue adapter that enforces per-queue concurrency limits, allowing users to control how many executors can work on each queue simultaneously (e.g. 3 for queue_a, 5 for queue_b). The adapter tracks in-flight job counts per queue and skips queues that have reached their ceiling. A new `release` hook on the base DequeueAdapter class is called by executors after job completion (and by the overseer on dead executor recovery) to free the slot. https://claude.ai/code/session_01FHk1mXeS9SDbwakWbKua4z
1 parent 6a96cc1 commit 31479ef

3 files changed

Lines changed: 239 additions & 0 deletions

File tree

Lines changed: 169 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,169 @@
1+
require "../../spec_helper"
2+
3+
describe "Mosquito::ConcurrencyLimitedDequeueAdapter" do
4+
getter(overseer : MockOverseer) { MockOverseer.new }
5+
getter(queue_list : MockQueueList) { overseer.queue_list.as(MockQueueList) }
6+
7+
def register(job_class : Mosquito::Job.class)
8+
Mosquito::Base.register_job_mapping job_class.name.underscore, job_class
9+
queue_list.queues << job_class.queue
10+
end
11+
12+
it "dequeues a job when under the limit" do
13+
clean_slate do
14+
register QueuedTestJob
15+
expected_job_run = QueuedTestJob.new.enqueue
16+
17+
adapter = Mosquito::ConcurrencyLimitedDequeueAdapter.new({
18+
"queued_test_job" => 3,
19+
})
20+
21+
result = adapter.dequeue(queue_list)
22+
refute_nil result
23+
if result
24+
assert_equal expected_job_run, result.job_run
25+
assert_equal QueuedTestJob.queue, result.queue
26+
end
27+
end
28+
end
29+
30+
it "returns nil when no jobs are available" do
31+
clean_slate do
32+
register QueuedTestJob
33+
34+
adapter = Mosquito::ConcurrencyLimitedDequeueAdapter.new({
35+
"queued_test_job" => 3,
36+
})
37+
38+
result = adapter.dequeue(queue_list)
39+
assert_nil result
40+
end
41+
end
42+
43+
it "skips a queue that has reached its concurrency limit" do
44+
clean_slate do
45+
register QueuedTestJob
46+
3.times { QueuedTestJob.new.enqueue }
47+
48+
adapter = Mosquito::ConcurrencyLimitedDequeueAdapter.new({
49+
"queued_test_job" => 2,
50+
})
51+
52+
# Dequeue twice — should succeed and fill the limit.
53+
result1 = adapter.dequeue(queue_list)
54+
refute_nil result1
55+
assert_equal 1, adapter.active_count("queued_test_job")
56+
57+
result2 = adapter.dequeue(queue_list)
58+
refute_nil result2
59+
assert_equal 2, adapter.active_count("queued_test_job")
60+
61+
# Third dequeue should be blocked by the limit.
62+
result3 = adapter.dequeue(queue_list)
63+
assert_nil result3
64+
end
65+
end
66+
67+
it "allows dequeue again after finished_with" do
68+
clean_slate do
69+
register QueuedTestJob
70+
3.times { QueuedTestJob.new.enqueue }
71+
72+
adapter = Mosquito::ConcurrencyLimitedDequeueAdapter.new({
73+
"queued_test_job" => 1,
74+
})
75+
76+
# Fill the single slot.
77+
result1 = adapter.dequeue(queue_list)
78+
refute_nil result1
79+
assert_equal 1, adapter.active_count("queued_test_job")
80+
81+
# Blocked.
82+
result2 = adapter.dequeue(queue_list)
83+
assert_nil result2
84+
85+
# Signal that the job finished.
86+
adapter.finished_with(result1.not_nil!.job_run, result1.not_nil!.queue)
87+
assert_equal 0, adapter.active_count("queued_test_job")
88+
89+
# Now dequeue should work again.
90+
result3 = adapter.dequeue(queue_list)
91+
refute_nil result3
92+
end
93+
end
94+
95+
it "does not limit queues not in the limits table" do
96+
clean_slate do
97+
register QueuedTestJob
98+
5.times { QueuedTestJob.new.enqueue }
99+
100+
# No limit configured for queued_test_job.
101+
adapter = Mosquito::ConcurrencyLimitedDequeueAdapter.new({
102+
"other_queue" => 1,
103+
})
104+
105+
# Should dequeue all 5 without blocking.
106+
5.times do |i|
107+
result = adapter.dequeue(queue_list)
108+
refute_nil result, "Expected dequeue ##{i + 1} to succeed"
109+
end
110+
end
111+
end
112+
113+
it "enforces independent limits across multiple queues" do
114+
clean_slate do
115+
register QueuedTestJob
116+
register EchoJob
117+
3.times { QueuedTestJob.new.enqueue }
118+
3.times { EchoJob.new(text: "hello").enqueue }
119+
120+
adapter = Mosquito::ConcurrencyLimitedDequeueAdapter.new({
121+
"queued_test_job" => 1,
122+
"io_queue" => 2,
123+
})
124+
125+
# Saturate queued_test_job (limit 1).
126+
# Because of shuffle we may get either queue first, so keep
127+
# dequeuing until the counters match the limits.
128+
results = [] of Mosquito::WorkUnit
129+
6.times do
130+
if r = adapter.dequeue(queue_list)
131+
results << r
132+
end
133+
end
134+
135+
assert_equal 1, adapter.active_count("queued_test_job")
136+
assert_equal 2, adapter.active_count("io_queue")
137+
assert_equal 3, results.size
138+
end
139+
end
140+
141+
it "finished_with does not go below zero" do
142+
adapter = Mosquito::ConcurrencyLimitedDequeueAdapter.new({
143+
"queued_test_job" => 3,
144+
})
145+
146+
job_run = Mosquito::JobRun.new("queued_test_job")
147+
queue = Mosquito::Queue.new("queued_test_job")
148+
adapter.finished_with(job_run, queue)
149+
assert_equal 0, adapter.active_count("queued_test_job")
150+
end
151+
152+
it "can be used via the overseer" do
153+
clean_slate do
154+
adapter = Mosquito::ConcurrencyLimitedDequeueAdapter.new({
155+
"queued_test_job" => 5,
156+
})
157+
overseer.dequeue_adapter = adapter
158+
159+
register QueuedTestJob
160+
expected_job_run = QueuedTestJob.new.enqueue
161+
162+
result = overseer.dequeue_job?
163+
refute_nil result
164+
if result
165+
assert_equal expected_job_run, result.job_run
166+
end
167+
end
168+
end
169+
end
Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
require "../dequeue_adapter"
2+
3+
module Mosquito
4+
# A dequeue adapter that enforces per-queue concurrency limits.
5+
#
6+
# Each queue can be assigned a maximum number of jobs that may execute
7+
# concurrently. When a queue has reached its limit, it is skipped during
8+
# dequeue until an in-flight job finishes.
9+
#
10+
# Queues not present in the limits table have no concurrency ceiling and
11+
# are bounded only by the total executor pool size.
12+
#
13+
# Among eligible queues the adapter uses a shuffle to provide rough
14+
# fairness, similar to `ShuffleDequeueAdapter`.
15+
#
16+
# ## Example
17+
#
18+
# ```crystal
19+
# Mosquito.configure do |settings|
20+
# settings.executor_count = 8
21+
#
22+
# settings.dequeue_adapter = Mosquito::ConcurrencyLimitedDequeueAdapter.new({
23+
# "queue_a" => 3,
24+
# "queue_b" => 5,
25+
# })
26+
# end
27+
# ```
28+
#
29+
# In this configuration at most 3 jobs from "queue_a" and 5 from "queue_b"
30+
# will execute at the same time. Other queues are unlimited.
31+
class ConcurrencyLimitedDequeueAdapter < DequeueAdapter
32+
getter limits : Hash(String, Int32)
33+
34+
# Tracks the number of currently in-flight jobs per queue name.
35+
# Access is fiber-safe because Crystal fibers are cooperatively
36+
# scheduled and we never yield between read and write.
37+
@active : Hash(String, Int32)
38+
39+
def initialize(@limits : Hash(String, Int32))
40+
@active = Hash(String, Int32).new(0)
41+
end
42+
43+
def dequeue(queue_list : Runners::QueueList) : WorkUnit?
44+
queue_list.queues.shuffle.each do |q|
45+
if limit = limits[q.name]?
46+
next if @active[q.name] >= limit
47+
end
48+
49+
if job_run = q.dequeue
50+
@active[q.name] = @active[q.name] + 1
51+
return WorkUnit.of(job_run, from: q)
52+
end
53+
end
54+
end
55+
56+
# Called by the Overseer when a job from this queue has finished
57+
# executing. Decrements the in-flight counter so the queue becomes
58+
# eligible for dequeue again.
59+
def finished_with(job_run : JobRun, queue : Queue) : Nil
60+
count = @active[queue.name]
61+
@active[queue.name] = {count - 1, 0}.max
62+
end
63+
64+
# Returns the current number of in-flight jobs for the given queue.
65+
def active_count(queue_name : String) : Int32
66+
@active[queue_name]
67+
end
68+
end
69+
end

src/mosquito/runners/overseer.cr

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -293,6 +293,7 @@ module Mosquito::Runners
293293
return unless work_unit = dead_executor.work_unit?
294294

295295
observer.recovered_job_from_executor work_unit.job_run, dead_executor
296+
dequeue_adapter.finished_with(work_unit.job_run, work_unit.queue)
296297
work_unit.job_run.retry_or_banish work_unit.queue
297298
end
298299

0 commit comments

Comments
 (0)