2 files changed, 38 insertions, 147 deletions

diff --git a/spec/workers/concerns/limited_capacity/job_tracker_spec.rb b/spec/workers/concerns/limited_capacity/job_tracker_spec.rb
index 2c79f347903..f141a1ad7ad 100644
--- a/spec/workers/concerns/limited_capacity/job_tracker_spec.rb
+++ b/spec/workers/concerns/limited_capacity/job_tracker_spec.rb
@@ -7,30 +7,30 @@ RSpec.describe LimitedCapacity::JobTracker, :clean_gitlab_redis_queues do
     described_class.new('namespace')
   end
 
+  let(:max_jids) { 10 }
+
   describe '#register' do
     it 'adds jid to the set' do
-      job_tracker.register('a-job-id')
-
+      expect(job_tracker.register('a-job-id', max_jids)). to be true
       expect(job_tracker.running_jids).to contain_exactly('a-job-id')
     end
 
-    it 'updates the counter' do
-      expect { job_tracker.register('a-job-id') }
-        .to change { job_tracker.count }
-        .from(0)
-        .to(1)
-    end
-
-    it 'does it in only one Redis call' do
-      expect(job_tracker).to receive(:with_redis).once.and_call_original
+    it 'returns false if the jid was not added' do
+      max_jids = 2
+      %w[jid1 jid2].each do |jid|
+        expect(job_tracker.register(jid, max_jids)).to be true
+      end
 
-      job_tracker.register('a-job-id')
+      expect(job_tracker.register('jid3', max_jids)).to be false
+      expect(job_tracker.running_jids).to contain_exactly(*%w[jid1 jid2])
     end
   end
 
   describe '#remove' do
     before do
-      job_tracker.register(%w[a-job-id other-job-id])
+      %w[a-job-id other-job-id].each do |jid|
+        job_tracker.register(jid, max_jids)
+      end
     end
 
     it 'removes jid from the set' do
@@ -38,24 +38,11 @@ RSpec.describe LimitedCapacity::JobTracker, :clean_gitlab_redis_queues do
 
       expect(job_tracker.running_jids).to contain_exactly('a-job-id')
     end
-
-    it 'updates the counter' do
-      expect { job_tracker.remove('other-job-id') }
-        .to change { job_tracker.count }
-        .from(2)
-        .to(1)
-    end
-
-    it 'does it in only one Redis call' do
-      expect(job_tracker).to receive(:with_redis).once.and_call_original
-
-      job_tracker.remove('other-job-id')
-    end
   end
 
   describe '#clean_up' do
     before do
-      job_tracker.register('a-job-id')
+      job_tracker.register('a-job-id', max_jids)
     end
 
     context 'with running jobs' do
@@ -83,13 +70,6 @@ RSpec.describe LimitedCapacity::JobTracker, :clean_gitlab_redis_queues do
           .to change { job_tracker.running_jids.include?('a-job-id') }
       end
 
-      it 'updates the counter' do
-        expect { job_tracker.clean_up }
-          .to change { job_tracker.count }
-          .from(1)
-          .to(0)
-      end
-
       it 'gets the job ids, removes them, and updates the counter with only two Redis calls' do
         expect(job_tracker).to receive(:with_redis).twice.and_call_original
 
diff --git a/spec/workers/concerns/limited_capacity/worker_spec.rb b/spec/workers/concerns/limited_capacity/worker_spec.rb
index 2c33c8666ec..790b5c3544d 100644
--- a/spec/workers/concerns/limited_capacity/worker_spec.rb
+++ b/spec/workers/concerns/limited_capacity/worker_spec.rb
@@ -44,40 +44,22 @@ RSpec.describe LimitedCapacity::Worker, :clean_gitlab_redis_queues, :aggregate_f
   describe '.perform_with_capacity' do
     subject(:perform_with_capacity) { worker_class.perform_with_capacity(:arg) }
 
+    let(:max_running_jobs) { 3 }
+
     before do
       expect_next_instance_of(worker_class) do |instance|
         expect(instance).to receive(:remove_failed_jobs)
-        expect(instance).to receive(:report_prometheus_metrics)
-
-        allow(instance).to receive(:remaining_work_count).and_return(remaining_work_count)
-        allow(instance).to receive(:remaining_capacity).and_return(remaining_capacity)
-      end
-    end
-
-    context 'when capacity is larger than work' do
-      let(:remaining_work_count) { 2 }
-      let(:remaining_capacity) { 3 }
 
-      it 'enqueues jobs for remaining work' do
-        expect(worker_class)
-          .to receive(:bulk_perform_async)
-          .with([[:arg], [:arg]])
-
-        perform_with_capacity
+        allow(instance).to receive(:max_running_jobs).and_return(max_running_jobs)
       end
     end
 
-    context 'when capacity is lower than work' do
-      let(:remaining_work_count) { 5 }
-      let(:remaining_capacity) { 3 }
-
-      it 'enqueues jobs for remaining work' do
-        expect(worker_class)
-          .to receive(:bulk_perform_async)
-          .with([[:arg], [:arg], [:arg]])
+    it 'enqueues jobs' do
+      expect(worker_class)
+        .to receive(:bulk_perform_async)
+        .with([[:arg], [:arg], [:arg]])
 
-        perform_with_capacity
-      end
+      perform_with_capacity
     end
   end
 
@@ -104,34 +86,27 @@ RSpec.describe LimitedCapacity::Worker, :clean_gitlab_redis_queues, :aggregate_f
         perform
       end
 
-      it 'registers itself in the running set' do
+      it 'reports prometheus metrics' do
         allow(worker).to receive(:perform_work)
-        expect(job_tracker).to receive(:register).with('my-jid')
+        expect(worker).to receive(:report_prometheus_metrics).once.and_call_original
+        expect(worker).to receive(:report_running_jobs_metrics).twice.and_call_original
 
         perform
       end
 
-      it 'removes itself from the running set' do
-        expect(job_tracker).to receive(:remove).with('my-jid')
-
+      it 'updates the running set' do
+        expect(job_tracker.running_jids).to be_empty
         allow(worker).to receive(:perform_work)
 
         perform
-      end
 
-      it 'reports prometheus metrics' do
-        allow(worker).to receive(:perform_work)
-        expect(worker).to receive(:report_prometheus_metrics).once.and_call_original
-        expect(worker).to receive(:report_running_jobs_metrics).twice.and_call_original
-
-        perform
+        expect(job_tracker.running_jids).to be_empty
       end
     end
 
     context 'with capacity and without work' do
       before do
         allow(worker).to receive(:max_running_jobs).and_return(10)
-        allow(worker).to receive(:running_jobs_count).and_return(0)
         allow(worker).to receive(:remaining_work_count).and_return(0)
         allow(worker).to receive(:perform_work)
       end
@@ -146,7 +121,7 @@ RSpec.describe LimitedCapacity::Worker, :clean_gitlab_redis_queues, :aggregate_f
     context 'without capacity' do
       before do
         allow(worker).to receive(:max_running_jobs).and_return(10)
-        allow(worker).to receive(:running_jobs_count).and_return(15)
+        allow(job_tracker).to receive(:register).and_return(false)
         allow(worker).to receive(:remaining_work_count).and_return(10)
       end
 
@@ -161,27 +136,14 @@ RSpec.describe LimitedCapacity::Worker, :clean_gitlab_redis_queues, :aggregate_f
 
         perform
       end
-
-      it 'does not register in the running set' do
-        expect(job_tracker).not_to receive(:register)
-
-        perform
-      end
-
-      it 'removes itself from the running set' do
-        expect(job_tracker).to receive(:remove).with('my-jid')
-
-        perform
-      end
-
-      it 'reports prometheus metrics' do
-        expect(worker).to receive(:report_prometheus_metrics)
-
-        perform
-      end
     end
 
     context 'when perform_work fails' do
+      before do
+        allow(worker).to receive(:max_running_jobs).and_return(10)
+        allow(job_tracker).to receive(:register).and_return(true)
+      end
+
       it 'does not re-enqueue itself' do
         expect(worker).not_to receive(:re_enqueue)
 
@@ -189,7 +151,7 @@ RSpec.describe LimitedCapacity::Worker, :clean_gitlab_redis_queues, :aggregate_f
       end
 
       it 'removes itself from the running set' do
-        expect(job_tracker).to receive(:remove)
+        expect(job_tracker).to receive(:remove).with('my-jid')
 
         expect { perform }.to raise_error(NotImplementedError)
       end
@@ -202,65 +164,14 @@ RSpec.describe LimitedCapacity::Worker, :clean_gitlab_redis_queues, :aggregate_f
     end
   end
 
-  describe '#remaining_capacity' do
-    subject(:remaining_capacity) { worker.remaining_capacity }
-
-    before do
-      expect(worker).to receive(:max_running_jobs).and_return(max_capacity)
-    end
-
-    context 'when changing the capacity to a lower value' do
-      let(:max_capacity) { -1 }
-
-      it { expect(remaining_capacity).to eq(0) }
-    end
-
-    context 'when registering new jobs' do
-      let(:max_capacity) { 2 }
-
-      before do
-        job_tracker.register('a-job-id')
-      end
-
-      it { expect(remaining_capacity).to eq(1) }
-    end
-
-    context 'with jobs in the queue' do
-      let(:max_capacity) { 2 }
-
-      before do
-        expect(worker_class).to receive(:queue_size).and_return(1)
-      end
-
-      it { expect(remaining_capacity).to eq(1) }
-    end
-
-    context 'with both running jobs and queued jobs' do
-      let(:max_capacity) { 10 }
-
-      before do
-        expect(worker_class).to receive(:queue_size).and_return(5)
-        expect(worker).to receive(:running_jobs_count).and_return(3)
-      end
-
-      it { expect(remaining_capacity).to eq(2) }
-    end
-  end
-
   describe '#remove_failed_jobs' do
     subject(:remove_failed_jobs) { worker.remove_failed_jobs }
 
-    before do
-      job_tracker.register('a-job-id')
-      allow(worker).to receive(:max_running_jobs).and_return(2)
+    it 'removes failed jobs' do
+      job_tracker.register('a-job-id', 10)
 
       expect(job_tracker).to receive(:clean_up).and_call_original
-    end
-
-    context 'with failed jobs' do
-      it 'update the available capacity' do
-        expect { remove_failed_jobs }.to change { worker.remaining_capacity }.by(1)
-      end
+      expect { remove_failed_jobs }.to change { job_tracker.running_jids.size }.by(-1)
     end
   end