diff --git a/app/workers/concerns/application_worker.rb b/app/workers/concerns/application_worker.rb index 62748808ff18ea5e1e50d0c6d7311f67f1b8ac54..733156ab758febaad5d0a0783169b25403df919a 100644 --- a/app/workers/concerns/application_worker.rb +++ b/app/workers/concerns/application_worker.rb @@ -9,6 +9,7 @@ module ApplicationWorker include Sidekiq::Worker # rubocop:disable Cop/IncludeSidekiqWorker include WorkerAttributes + include WorkerContext included do set_queue diff --git a/app/workers/concerns/cronjob_queue.rb b/app/workers/concerns/cronjob_queue.rb index 0683b2293815f46a25e0491cd1be2f7f9f5ab34a..25ee4539cab10bac17378afc65d09da4774da8e4 100644 --- a/app/workers/concerns/cronjob_queue.rb +++ b/app/workers/concerns/cronjob_queue.rb @@ -8,5 +8,6 @@ module CronjobQueue included do queue_namespace :cronjob sidekiq_options retry: false + worker_context project: nil, namespace: nil, user: nil end end diff --git a/app/workers/concerns/worker_context.rb b/app/workers/concerns/worker_context.rb new file mode 100644 index 0000000000000000000000000000000000000000..a497c720cfa97578ead53db36b570252f45cb1c2 --- /dev/null +++ b/app/workers/concerns/worker_context.rb @@ -0,0 +1,52 @@ +# frozen_string_literal: true + +module WorkerContext + extend ActiveSupport::Concern + + class_methods do + def bulk_perform_async_with_contexts(contexts_for_arguments) + with_batch_contexts(contexts_for_arguments) do + bulk_perform_async(contexts_for_arguments.keys) + end + end + + def bulk_perform_in_with_contexts(delay, contexts_for_arguments) + with_batch_contexts(contexts_for_arguments) do + bulk_perform_in(delay, contexts_for_arguments.keys) + end + end + + def worker_context(attributes) + @worker_context = Gitlab::ApplicationContext.new(attributes) + end + + def get_worker_context + @worker_context + end + + def context_for_arguments(args) + batch_scheduling_contexts[args] + end + + private + + def batch_scheduling_contexts + Thread.current["#{name}_batch_scheduling_contexts"] ||= {} + end + + def with_batch_contexts(contexts_for_atributes) + contexts = contexts_for_atributes.transform_values do |info| + Gitlab::ApplicationContext.new(info) + end + batch_scheduling_contexts.merge!(contexts) + + yield + ensure + batch_scheduling_contexts.clear + end + end + + def with_context(context, &block) + Gitlab::ApplicationContext.new(context).use(&block) + end +end diff --git a/doc/development/sidekiq_style_guide.md b/doc/development/sidekiq_style_guide.md index 77663b0bb29f6ada9f1375636b1272c97779b83f..0fc6033ea8b48459ef9119c8bb359cdfde37093c 100644 --- a/doc/development/sidekiq_style_guide.md +++ b/doc/development/sidekiq_style_guide.md @@ -273,6 +273,10 @@ class SomeCrossCuttingConcernWorker end ``` +## Worker context + +### Cronjob workers + ## Tests Each Sidekiq worker must be tested using RSpec, just like any other class. These diff --git a/ee/app/workers/repository_update_mirror_worker.rb b/ee/app/workers/repository_update_mirror_worker.rb index 4436fad6043dfcadcb9705433828d49637a5d405..bc624b113c20ae0f5fa76b68578140bbf125cc5f 100644 --- a/ee/app/workers/repository_update_mirror_worker.rb +++ b/ee/app/workers/repository_update_mirror_worker.rb @@ -20,11 +20,12 @@ def perform(project_id) return unless start_mirror(project) @current_user = project.mirror_user || project.creator + with_context(project: project, user: @current_user) do + result = Projects::UpdateMirrorService.new(project, @current_user).execute + raise UpdateError, result[:message] if result[:status] == :error - result = Projects::UpdateMirrorService.new(project, @current_user).execute - raise UpdateError, result[:message] if result[:status] == :error - - finish_mirror(project) + finish_mirror(project) + end rescue UpdateError => ex fail_mirror(project, ex.message) raise diff --git a/ee/app/workers/update_all_mirrors_worker.rb b/ee/app/workers/update_all_mirrors_worker.rb index 13cf8d17b9a21909168783191d59b16efa38ef37..1b6dc923980c6aaaef4e46638ffdf74c7dd296d2 100644 --- a/ee/app/workers/update_all_mirrors_worker.rb +++ b/ee/app/workers/update_all_mirrors_worker.rb @@ -14,24 +14,22 @@ class UpdateAllMirrorsWorker def perform return if Gitlab::Database.read_only? - Gitlab::ApplicationContext.with_context({ user: nil, project: nil, namespace: nil }) do - scheduled = 0 - with_lease do - scheduled = schedule_mirrors! - end - - # If we didn't get the lease, or no updates were scheduled, exit early - break unless scheduled > 0 - - # Wait to give some jobs a chance to complete - Kernel.sleep(RESCHEDULE_WAIT) - - # If there's capacity left now (some jobs completed), - # reschedule this job to enqueue more work. - # - # This is in addition to the regular (cron-like) scheduling of this job. - UpdateAllMirrorsWorker.perform_async if Gitlab::Mirror.reschedule_immediately? + scheduled = 0 + with_lease do + scheduled = schedule_mirrors! end + + # If we didn't get the lease, or no updates were scheduled, exit early + return unless scheduled > 0 + + # Wait to give some jobs a chance to complete + Kernel.sleep(RESCHEDULE_WAIT) + + # If there's capacity left now (some jobs completed), + # reschedule this job to enqueue more work. + # + # This is in addition to the regular (cron-like) scheduling of this job. + UpdateAllMirrorsWorker.perform_async if Gitlab::Mirror.reschedule_immediately? end # rubocop: disable CodeReuse/ActiveRecord @@ -49,11 +47,12 @@ def schedule_mirrors! projects = pull_mirrors_batch(freeze_at: now, batch_size: batch_size, offset_at: last).to_a break if projects.empty? - project_ids = projects.lazy.select(&:mirror?).take(capacity).map(&:id).force - capacity -= project_ids.length + projects_to_schedule = projects.lazy.select(&:mirror?).take(capacity).force + capacity -= projects_to_schedule.size + + schedule_projects_in_batch(projects_to_schedule) - ProjectImportScheduleWorker.bulk_perform_async(project_ids.map { |id| [id] }) - scheduled += project_ids.length + scheduled += projects_to_schedule.length # If fewer than `batch_size` projects were returned, we don't need to query again break if projects.length < batch_size @@ -97,6 +96,7 @@ def pull_mirrors_batch(freeze_at:, batch_size:, offset_at: nil) .mirrors_to_sync(freeze_at) .reorder('import_state.next_execution_timestamp') .limit(batch_size) + .with_route .includes(:namespace) # Used by `project.mirror?` relation = relation.where('import_state.next_execution_timestamp > ?', offset_at) if offset_at @@ -104,4 +104,13 @@ def pull_mirrors_batch(freeze_at:, batch_size:, offset_at: nil) relation end # rubocop: enable CodeReuse/ActiveRecord + + def schedule_projects_in_batch(projects) + contexts_for_arguments = {} + projects.each do |project| + contexts_for_arguments[[project.id]] = { project: project } + end + + ProjectImportScheduleWorker.bulk_perform_async_with_contexts(contexts_for_arguments) + end end diff --git a/lib/gitlab/application_context.rb b/lib/gitlab/application_context.rb index 71dbfea70e8e6b9039049caa9177a5d8a5e8d82e..80b6b74d8c3afc22278e13834febf3ac4fd9b784 100644 --- a/lib/gitlab/application_context.rb +++ b/lib/gitlab/application_context.rb @@ -15,7 +15,7 @@ class ApplicationContext def self.with_context(args, &block) application_context = new(**args) - Labkit::Context.with_context(application_context.to_lazy_hash, &block) + application_context.use(&block) end def self.push(args) @@ -40,6 +40,10 @@ def to_lazy_hash end end + def use + Labkit::Context.with_context(to_lazy_hash) { yield } + end + private attr_reader :set_values diff --git a/lib/gitlab/sidekiq_middleware.rb b/lib/gitlab/sidekiq_middleware.rb index 3dda244233f5df81f9a6fda2edb74edabf995388..d943ce92e4c089047dc251740465667d52a4e5ba 100644 --- a/lib/gitlab/sidekiq_middleware.rb +++ b/lib/gitlab/sidekiq_middleware.rb @@ -18,6 +18,7 @@ def self.server_configurator(metrics: true, arguments_logger: true, memory_kille chain.add Labkit::Middleware::Sidekiq::Server chain.add Gitlab::SidekiqMiddleware::InstrumentationLogger chain.add Gitlab::SidekiqStatus::ServerMiddleware + chain.add Gitlab::SidekiqMiddleware::WorkerContext::Server end end @@ -28,6 +29,7 @@ def self.client_configurator lambda do |chain| chain.add Gitlab::SidekiqStatus::ClientMiddleware chain.add Gitlab::SidekiqMiddleware::ClientMetrics + chain.add Gitlab::SidekiqMiddleware::WorkerContext::Client chain.add Labkit::Middleware::Sidekiq::Client end end diff --git a/lib/gitlab/sidekiq_middleware/worker_context.rb b/lib/gitlab/sidekiq_middleware/worker_context.rb new file mode 100644 index 0000000000000000000000000000000000000000..cf59a5de33d1e296154e2644332aed7d7c9d4bc6 --- /dev/null +++ b/lib/gitlab/sidekiq_middleware/worker_context.rb @@ -0,0 +1,29 @@ +# frozen_string_literal: true + +module Gitlab + module SidekiqMiddleware + module WorkerContext + class Server + include SidekiqMiddleware::WorkerContext + + def call(worker, job, _queue, &block) + worker_class = worker.class + + # This is not a worker we know about, perhaps from a gem + return yield unless worker_class.include?(ApplicationWorker) + + # Use the context defined on the class level as a base context + wrap_in_context(worker_class.get_worker_context, &block) + end + end + + private + + def wrap_in_optional_context(context_or_nil, &block) + return yield unless context_or_nil + + context_or_nil.use(&block) + end + end + end +end diff --git a/lib/gitlab/sidekiq_middleware/worker_context/client.rb b/lib/gitlab/sidekiq_middleware/worker_context/client.rb new file mode 100644 index 0000000000000000000000000000000000000000..dabd14ca9bec96d1f17b4b8884f13368c500f1a1 --- /dev/null +++ b/lib/gitlab/sidekiq_middleware/worker_context/client.rb @@ -0,0 +1,29 @@ +# frozen_string_literal: true + +module Gitlab + module SidekiqMiddleware + module WorkerContext + class Client + include Gitlab::SidekiqMiddleware::WorkerContext + + def call(worker_class_or_name, job, _queue, _redis_pool, &block) + worker_class = constantize_worker(worker_class_or_name) + # Mailers can't be constantized like this + return yield unless worker_class + + context_for_args = worker_class.context_for_arguments(job['args']) + + wrap_in_optional_context(context_for_args, &block) + end + + private + + def constantize_worker(class_name) + class_name.to_s.constantize + rescue NameError + nil + end + end + end + end +end diff --git a/lib/gitlab/sidekiq_middleware/worker_context/server.rb b/lib/gitlab/sidekiq_middleware/worker_context/server.rb new file mode 100644 index 0000000000000000000000000000000000000000..ae9d1ef3040974702e4d2e7461d5e2dacf07352a --- /dev/null +++ b/lib/gitlab/sidekiq_middleware/worker_context/server.rb @@ -0,0 +1,21 @@ +# frozen_string_literal: true + +module Gitlab + module SidekiqMiddleware + module WorkerContext + class Server + include SidekiqMiddleware::WorkerContext + + def call(worker, job, _queue, &block) + worker_class = worker.class + + # This is not a worker we know about, perhaps from a gem + return yield unless worker_class.include?(ApplicationWorker) + + # Use the context defined on the class level as a base context + wrap_in_optional_context(worker_class.get_worker_context, &block) + end + end + end + end +end diff --git a/rubocop/cop/scalability/bulk_perform_with_context.rb b/rubocop/cop/scalability/bulk_perform_with_context.rb new file mode 100644 index 0000000000000000000000000000000000000000..6118d4254f7d72bdf195d2ebed29cbe1fd5c0a0b --- /dev/null +++ b/rubocop/cop/scalability/bulk_perform_with_context.rb @@ -0,0 +1,28 @@ +# frozen_string_literal: true + +module RuboCop + module Cop + module Scalability + class BulkPerformWithContext < RuboCop::Cop::Cop + MSG = <<~MSG.freeze + Prefer using `Worker.bulk_perform_async_with_contexts` and + `Worker.bulk_perform_in_with_context` over the methods without a context + if your worker deals with specific projects or namespaces + The context is required to add metadata to our logs. + + Read more about it https://docs.gitlab.com/ee/development/sidekiq_style_guide.html#worker-context + MSG + + def_node_matcher :schedules_in_batch_without_context?, <<~PATTERN + (send (...) {:bulk_perform_async :bulk_perform_in} (...)) + PATTERN + + def on_send(node) + return unless schedules_in_batch_without_context?(node) + + add_offense(node, location: :expression) + end + end + end + end +end diff --git a/rubocop/cop/scalability/cron_worker_context.rb b/rubocop/cop/scalability/cron_worker_context.rb new file mode 100644 index 0000000000000000000000000000000000000000..d057690019cedc23be5f30461b61f099af345b44 --- /dev/null +++ b/rubocop/cop/scalability/cron_worker_context.rb @@ -0,0 +1,36 @@ +# frozen_string_literal: true + +module RuboCop + module Cop + module Scalability + class CronWorkerContext < RuboCop::Cop::Cop + MSG = <<~MSG.freeze + Manually define an ApplicationContext for cronjob-workers. The context + is required to add metadata to our logs. + + Read more about it https://docs.gitlab.com/ee/development/sidekiq_style_guide.html#worker-context + MSG + + def_node_matcher :includes_cronjob_queue?, <<~PATTERN + (send nil? :include (const nil? :CronjobQueue)) + PATTERN + + def_node_search :defines_contexts?, <<~PATTERN + (send nil? :with_context _) + PATTERN + + def_node_search :schedules_with_batch_context?, <<~PATTERN + (send (...) {:bulk_perform_async_with_contexts :bulk_perform_in_with_contexts} (...)) + PATTERN + + def on_send(node) + return unless includes_cronjob_queue?(node) + return if defines_contexts?(node.parent) + return if schedules_with_batch_context?(node.parent) + + add_offense(node.arguments.first, location: :expression) + end + end + end + end +end diff --git a/rubocop/rubocop.rb b/rubocop/rubocop.rb index 1479dc3384a461935957037264596cf58bd6cbc9..6410129fcad4698540f3b4490e50dc54910da46c 100644 --- a/rubocop/rubocop.rb +++ b/rubocop/rubocop.rb @@ -45,6 +45,8 @@ require_relative 'cop/qa/ambiguous_page_object_name' require_relative 'cop/sidekiq_options_queue' require_relative 'cop/scalability/file_uploads' +require_relative 'cop/scalability/bulk_perform_with_context' +require_relative 'cop/scalability/cron_worker_context' require_relative 'cop/destroy_all' require_relative 'cop/ruby_interpolation_in_translation' require_relative 'code_reuse_helpers' diff --git a/spec/lib/gitlab/application_context_spec.rb b/spec/lib/gitlab/application_context_spec.rb index 482bf0dc192d117afb6d7be4cf16bc9cf9a25e1e..1918cbf6e4508301d4703bfb0c4b424a4aaede05 100644 --- a/spec/lib/gitlab/application_context_spec.rb +++ b/spec/lib/gitlab/application_context_spec.rb @@ -101,4 +101,18 @@ def result(context) end end end + + describe '#use' do + let(:context) { described_class.new(user: build(:user)) } + + it 'yields control' do + expect { |b| context.use(&b) }.to yield_control + end + + it 'passes the expected context on to labkit' do + expect(Labkit::Context).to receive(:with_context).with(a_hash_including(user: duck_type(:call))) + + context.use {} + end + end end diff --git a/spec/lib/gitlab/sidekiq_middleware/worker_context/client_spec.rb b/spec/lib/gitlab/sidekiq_middleware/worker_context/client_spec.rb new file mode 100644 index 0000000000000000000000000000000000000000..20fb2041cf4f6e64264855c304a2fb1cfea7071c --- /dev/null +++ b/spec/lib/gitlab/sidekiq_middleware/worker_context/client_spec.rb @@ -0,0 +1,34 @@ +# frozen_string_literal: true + +require 'spec_helper' + +describe Gitlab::SidekiqMiddleware::WorkerContext::Client do + class TestWorker + include ApplicationWorker + + def self.job_for_args(args) + jobs.find { |job| job['args'] == args } + end + + def perform(*args) + end + end + + describe "#call" do + it 'applies a context for jobs scheduled in batch' do + user_1 = build_stubbed(:user, username: "user-1") + user_2 = build_stubbed(:user, username: "user-2") + + TestWorker.bulk_perform_async_with_contexts( + ['job1', 1, 2, 3] => { user: user_1 }, + ['job2', 1, 2, 3] => { user: user_2 } + ) + + job1 = TestWorker.job_for_args(['job1', 1, 2, 3]) + job2 = TestWorker.job_for_args(['job2', 1, 2, 3]) + + expect(job1['meta.user']).to eq(user_1.username) + expect(job2['meta.user']).to eq(user_2.username) + end + end +end diff --git a/spec/lib/gitlab/sidekiq_middleware/worker_context/server_spec.rb b/spec/lib/gitlab/sidekiq_middleware/worker_context/server_spec.rb new file mode 100644 index 0000000000000000000000000000000000000000..8fbcf98db586fceecabae8638b672f11ff2c427c --- /dev/null +++ b/spec/lib/gitlab/sidekiq_middleware/worker_context/server_spec.rb @@ -0,0 +1,48 @@ +# frozen_string_literal: true + +require 'spec_helper' + +describe Gitlab::SidekiqMiddleware::WorkerContext::Server do + class TestWorker + # To keep track of the context that was active for certain arguments + cattr_accessor(:contexts) { {} } + + include ApplicationWorker + + worker_context user: nil + + def perform(identifier, *args) + self.class.contexts.merge!(identifier => Labkit::Context.current.to_h) + end + end + + before do + TestWorker.contexts.clear + end + + around do |example| + Sidekiq::Testing.inline! { example.run } + end + + before(:context) do + Sidekiq::Testing.server_middleware do |chain| + chain.add described_class + end + end + + after(:context) do + Sidekiq::Testing.server_middleware do |chain| + chain.remove described_class + end + end + + describe "#call" do + it 'applies a class context' do + Gitlab::ApplicationContext.with_context(user: build_stubbed(:user)) do + TestWorker.perform_async("identifier", 1) + end + + expect(TestWorker.contexts['identifier'].keys).not_to include('meta.user') + end + end +end diff --git a/spec/lib/gitlab/sidekiq_middleware_spec.rb b/spec/lib/gitlab/sidekiq_middleware_spec.rb index 473d85c0143dea77590faf13d677e6d34e148a11..8c93e4ce1f22085f8fe657935a74a67449e351f9 100644 --- a/spec/lib/gitlab/sidekiq_middleware_spec.rb +++ b/spec/lib/gitlab/sidekiq_middleware_spec.rb @@ -120,6 +120,7 @@ def perform(_arg) # This test ensures that this does not happen it "invokes the chain" do expect_any_instance_of(Gitlab::SidekiqStatus::ClientMiddleware).to receive(:call).with(*middleware_expected_args).once.and_call_original + expect_any_instance_of(Gitlab::SidekiqMiddleware::WorkerContext).to receive(:call).with(*middleware_expected_args).once.and_call_original expect_any_instance_of(Labkit::Middleware::Sidekiq::Client).to receive(:call).with(*middleware_expected_args).once.and_call_original expect { |b| chain.invoke(worker_class_arg, job, queue, redis_pool, &b) }.to yield_control.once diff --git a/spec/rubocop/cop/scalability/bulk_perform_with_context_spec.rb b/spec/rubocop/cop/scalability/bulk_perform_with_context_spec.rb new file mode 100644 index 0000000000000000000000000000000000000000..a4fef63bb4c2ef12a1f83bf913ce425b418dc6c4 --- /dev/null +++ b/spec/rubocop/cop/scalability/bulk_perform_with_context_spec.rb @@ -0,0 +1,29 @@ +# frozen_string_literal: true + +require 'fast_spec_helper' +require 'rubocop' +require_relative '../../../support/helpers/expect_offense' +require_relative '../../../../rubocop/cop/scalability/bulk_perform_with_context' + +describe RuboCop::Cop::Scalability::BulkPerformWithContext do + include CopHelper + include ExpectOffense + + subject(:cop) { described_class.new } + + it "adds an offense when calling bulk_perform_async" do + inspect_source(<<~CODE.strip_indent) + Worker.bulk_perform_async(args) + CODE + + expect(cop.offenses.size).to eq(1) + end + + it "adds an offense when calling bulk_perform_in" do + inspect_source(<<~CODE.strip_indent) + Worker.bulk_perform_in(args) + CODE + + expect(cop.offenses.size).to eq(1) + end +end diff --git a/spec/rubocop/cop/scalability/cron_worker_context_spec.rb b/spec/rubocop/cop/scalability/cron_worker_context_spec.rb new file mode 100644 index 0000000000000000000000000000000000000000..bf10b8dc02ca5fae66283b783df04591cbaea9f0 --- /dev/null +++ b/spec/rubocop/cop/scalability/cron_worker_context_spec.rb @@ -0,0 +1,78 @@ +# frozen_string_literal: true + +require 'fast_spec_helper' +require 'rubocop' +require_relative '../../../support/helpers/expect_offense' +require_relative '../../../../rubocop/cop/scalability/cron_worker_context' + +describe RuboCop::Cop::Scalability::CronWorkerContext do + include CopHelper + include ExpectOffense + + subject(:cop) { described_class.new } + + it 'adds an offense when including CronjobQueue' do + inspect_source(<<~CODE.strip_indent) + class SomeWorker + include CronjobQueue + end + CODE + + expect(cop.offenses.size).to eq(1) + end + + it 'does not add offenses for other workers' do + expect_no_offenses(<<~CODE.strip_indent) + class SomeWorker + end + CODE + end + + it 'does not add an offense when the class defines a context' do + expect_no_offenses(<<~CODE.strip_indent) + class SomeWorker + include CronjobQueue + + with_context user: 'bla' + end + CODE + end + + it 'does not add an offense when the worker calls `with_context`' do + expect_no_offenses(<<~CODE.strip_indent) + class SomeWorker + include CronjobQueue + + def perform + with_context(user: 'bla') do + # more work + end + end + end + CODE + end + + it 'does not add an offense when the worker calls `bulk_perform_async_with_contexts`' do + expect_no_offenses(<<~CODE.strip_indent) + class SomeWorker + include CronjobQueue + + def perform + SomeOtherWorker.bulk_perform_async_with_contexts(contexts_for_arguments) + end + end + CODE + end + + it 'does not add an offense when the worker calls `bulk_perform_in_with_contexts`' do + expect_no_offenses(<<~CODE.strip_indent) + class SomeWorker + include CronjobQueue + + def perform + SomeOtherWorker.bulk_perform_in_with_contexts(contexts_for_arguments) + end + end + CODE + end +end diff --git a/spec/workers/concerns/cronjob_queue_spec.rb b/spec/workers/concerns/cronjob_queue_spec.rb index cf4d47b7500848af8656a027fcea92b5208d0eb2..ea3b7bad2e1cb01b6bcb20b8d2732f4524203956 100644 --- a/spec/workers/concerns/cronjob_queue_spec.rb +++ b/spec/workers/concerns/cronjob_queue_spec.rb @@ -10,7 +10,7 @@ def self.name end include ApplicationWorker - include CronjobQueue + include CronjobQueue # rubocop:disable Scalability/CronWorkerContext end end @@ -21,4 +21,12 @@ def self.name it 'disables retrying of failed jobs' do expect(worker.sidekiq_options['retry']).to eq(false) end + + it 'automatically clears project, user and namespace from the context', :aggregate_failues do + worker_context = worker.get_worker_context.to_lazy_hash.transform_values(&:call) + + expect(worker_context[:user]).to be_nil + expect(worker_context[:root_namespace]).to be_nil + expect(worker_context[:project]).to be_nil + end end diff --git a/spec/workers/concerns/worker_context_spec.rb b/spec/workers/concerns/worker_context_spec.rb new file mode 100644 index 0000000000000000000000000000000000000000..2fd7509859b8c7bd23d53f9cd475ec5db79c07bf --- /dev/null +++ b/spec/workers/concerns/worker_context_spec.rb @@ -0,0 +1,80 @@ +# frozen_string_literal: true + +require 'spec_helper' + +describe WorkerContext do + let(:worker) do + Class.new do + def self.name + 'Gitlab::Foo::Bar::DummyWorker' + end + + include ApplicationWorker # and thus also `WorkerContext`. + end + end + + shared_examples 'tracking bulk scheduling contexts' do + let(:arguments_with_contexts) do + worker.__send__(:batch_scheduling_contexts) + end + + it 'keeps track of the context per key to schedule' do + # stub clearing the contexts, so we can check what's inside + allow(arguments_with_contexts).to receive(:clear) + + subject + + expect(worker.context_for_arguments(["hello"])).to be_a(Gitlab::ApplicationContext) + end + + it 'clears the contexts' do + subject + + expect(arguments_with_contexts).to be_empty + end + end + + describe '.bulk_perform_async_with_contexts' do + subject do + worker.bulk_perform_async_with_contexts(["hello"] => { user: build_stubbed(:user) }, ["world"] => {}) + end + + it 'calls bulk_perform_async with the arguments' do + expect(worker).to receive(:bulk_perform_async).with([["hello"], ["world"]]) + + subject + end + + it_behaves_like 'tracking bulk scheduling contexts' + end + + describe '.bulk_perform_in_with_contexts' do + subject do + worker.bulk_perform_in_with_contexts(10.minutes, ["hello"] => { user: build_stubbed(:user) }, ["world"] => {}) + end + + it 'calls bulk_perform_in with the arguments and delay' do + expect(worker).to receive(:bulk_perform_in).with(10.minutes, [["hello"], ["world"]]) + + subject + end + + it_behaves_like 'tracking bulk scheduling contexts' + end + + describe '.worker_context' do + it 'allows modifying the context for the entire worker' do + worker.worker_context(user: build_stubbed(:user)) + + expect(worker.get_worker_context).to be_a(Gitlab::ApplicationContext) + end + end + + describe '#with_context' do + it 'allows modifying context when the job is running' do + worker.new.with_context(user: build_stubbed(:user, username: 'jane-doe')) do + expect(Labkit::Context.current.to_h).to include('meta.user' => 'jane-doe') + end + end + end +end