From 09d0a34bc95bc87e77aa6eabc67b4cf86f0f19bc Mon Sep 17 00:00:00 2001
From: Bob Van Landuyt <bob@vanlanduyt.co>
Date: Thu, 9 Jan 2020 15:59:47 +0100
Subject: [PATCH] WIP: Allow workers to specify a context

This would allow us to add context information to workers based on the
arguments passed into the worker.
---
 app/workers/concerns/application_worker.rb    |  1 +
 app/workers/concerns/cronjob_queue.rb         |  8 +++++
 app/workers/concerns/worker_context.rb        | 21 ++++++++++++
 .../workers/project_import_schedule_worker.rb |  4 +++
 .../repository_update_mirror_worker.rb        |  7 +++-
 ee/app/workers/update_all_mirrors_worker.rb   | 32 +++++++++----------
 lib/gitlab/application_context.rb             |  6 +++-
 lib/gitlab/sidekiq_middleware.rb              |  1 +
 .../worker_context_middleware.rb              | 17 ++++++++++
 9 files changed, 78 insertions(+), 19 deletions(-)
 create mode 100644 app/workers/concerns/worker_context.rb
 create mode 100644 lib/gitlab/sidekiq_middleware/worker_context_middleware.rb

diff --git a/app/workers/concerns/application_worker.rb b/app/workers/concerns/application_worker.rb
index 62748808ff18ea..733156ab758feb 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 0683b2293815f4..7bfa3a08844b8b 100644
--- a/app/workers/concerns/cronjob_queue.rb
+++ b/app/workers/concerns/cronjob_queue.rb
@@ -8,5 +8,13 @@ module CronjobQueue
   included do
     queue_namespace :cronjob
     sidekiq_options retry: false
+
+    # Start the Cronjob worker with an empty context. This can be overridden in
+    # the worker implementing this module.
+    worker_context do |*args|
+      # TODO, we could add that the job is triggered by a cron here, I think that
+      # would belong in the `caller_id` that we're still building
+      Gitlab::ApplicationContext.new(user: nil, namespace: nil, project: nil)
+    end
   end
 end
diff --git a/app/workers/concerns/worker_context.rb b/app/workers/concerns/worker_context.rb
new file mode 100644
index 00000000000000..00b14494677f0e
--- /dev/null
+++ b/app/workers/concerns/worker_context.rb
@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module WorkerContext
+  extend ActiveSupport::Concern
+
+  class_methods do
+    def worker_context(&block)
+      @context_builder = block
+    end
+
+    def with_worker_context(*args, &block)
+      return yield unless context_builder
+
+      context_builder.call(*args).use(&block)
+    end
+
+    private
+
+    attr_accessor :context_builder
+  end
+end
diff --git a/ee/app/workers/project_import_schedule_worker.rb b/ee/app/workers/project_import_schedule_worker.rb
index 7c795fd4e64627..7908be362e600c 100644
--- a/ee/app/workers/project_import_schedule_worker.rb
+++ b/ee/app/workers/project_import_schedule_worker.rb
@@ -9,6 +9,10 @@ class ProjectImportScheduleWorker
   feature_category :importers
   sidekiq_options retry: false
 
+  worker_context do |project_id|
+    Gitlab::ApplicationContext.new(project: -> { Project.with_route.find(project_id) })
+  end
+
   # rubocop: disable CodeReuse/ActiveRecord
   def perform(project_id)
     return if Gitlab::Database.read_only?
diff --git a/ee/app/workers/repository_update_mirror_worker.rb b/ee/app/workers/repository_update_mirror_worker.rb
index 4436fad6043dfc..9827a2b23e94a3 100644
--- a/ee/app/workers/repository_update_mirror_worker.rb
+++ b/ee/app/workers/repository_update_mirror_worker.rb
@@ -12,7 +12,12 @@ class RepositoryUpdateMirrorWorker
   # Retry not necessary. It will try again at the next update interval.
   sidekiq_options retry: false, status_expiration: StuckImportJobsWorker::IMPORT_JOBS_EXPIRATION
 
-  attr_accessor :project, :repository, :current_user
+  worker_context do |project_id|
+    project = Project.with_route.find(project_id)
+    user = project.mirror_user || project.creator
+
+    Gitlab::ApplicationContext.new(user: user, project: project)
+  end
 
   def perform(project_id)
     project = Project.find(project_id)
diff --git a/ee/app/workers/update_all_mirrors_worker.rb b/ee/app/workers/update_all_mirrors_worker.rb
index 13cf8d17b9a219..dfb3e6b2da076b 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
diff --git a/lib/gitlab/application_context.rb b/lib/gitlab/application_context.rb
index b9190b519a0cf6..6e178962c07235 100644
--- a/lib/gitlab/application_context.rb
+++ b/lib/gitlab/application_context.rb
@@ -7,7 +7,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)
@@ -25,6 +25,10 @@ def to_lazy_hash
         root_namespace: -> { root_namespace_path } }
     end
 
+    def use(&block)
+      Labkit::Context.with_context(to_lazy_hash, &block)
+    end
+
     private
 
     lazy_attr_reader :user, type: User
diff --git a/lib/gitlab/sidekiq_middleware.rb b/lib/gitlab/sidekiq_middleware.rb
index 4893cbc1f45f7a..f3b31e25affee5 100644
--- a/lib/gitlab/sidekiq_middleware.rb
+++ b/lib/gitlab/sidekiq_middleware.rb
@@ -27,6 +27,7 @@ def self.server_configurator(metrics: true, arguments_logger: true, memory_kille
     def self.client_configurator
       lambda do |chain|
         chain.add Gitlab::SidekiqStatus::ClientMiddleware
+        chain.add Gitlab::SidekiqMiddleware::WorkerContextMiddleware # needs to be before labkit, so labkit gets the right context
         chain.add Labkit::Middleware::Sidekiq::Client
       end
     end
diff --git a/lib/gitlab/sidekiq_middleware/worker_context_middleware.rb b/lib/gitlab/sidekiq_middleware/worker_context_middleware.rb
new file mode 100644
index 00000000000000..27fc1553f69b58
--- /dev/null
+++ b/lib/gitlab/sidekiq_middleware/worker_context_middleware.rb
@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module SidekiqMiddleware
+    class WorkerContextMiddleware
+      def call(worker_class, job, _queue, _redis_pool)
+        klazz = worker_class.constantize
+
+        return yield unless klazz.respond_to?(:with_worker_context)
+
+        klazz.with_worker_context(*job['args']) do
+          yield
+        end
+      end
+    end
+  end
+end
-- 
GitLab