Add latest changes from gitlab-org/gitlab@master

3 files changed, 82 insertions, 0 deletions

diff --git a/lib/gitlab/marginalia.rb b/lib/gitlab/marginalia.rb
new file mode 100644
index 00000000000..2be96cecae3
--- /dev/null
+++ b/lib/gitlab/marginalia.rb
@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Marginalia
+    MARGINALIA_FEATURE_FLAG = :marginalia
+
+    def self.set_application_name
+      ::Marginalia.application_name = Gitlab.process_name
+    end
+
+    def self.enable_sidekiq_instrumentation
+      if Sidekiq.server?
+        ::Marginalia::SidekiqInstrumentation.enable!
+      end
+    end
+
+    def self.cached_feature_enabled?
+      !!@enabled
+    end
+
+    def self.set_feature_cache
+      # During db:create and db:bootstrap skip feature query as DB is not available yet.
+      return false unless ActiveRecord::Base.connected? && Gitlab::Database.cached_table_exists?('features')
+
+      @enabled = Feature.enabled?(MARGINALIA_FEATURE_FLAG)
+    end
+  end
+end
diff --git a/lib/gitlab/marginalia/active_record_instrumentation.rb b/lib/gitlab/marginalia/active_record_instrumentation.rb
new file mode 100644
index 00000000000..3266b9f8336
--- /dev/null
+++ b/lib/gitlab/marginalia/active_record_instrumentation.rb
@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+# Patch to annotate sql only when the feature is enabled.
+module Gitlab
+  module Marginalia
+    module ActiveRecordInstrumentation
+      def annotate_sql(sql)
+        Gitlab::Marginalia.cached_feature_enabled? ? super(sql) : sql
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/marginalia/comment.rb b/lib/gitlab/marginalia/comment.rb
new file mode 100644
index 00000000000..a0eee823763
--- /dev/null
+++ b/lib/gitlab/marginalia/comment.rb
@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+# Module to support correlation_id and additional job details.
+module Gitlab
+  module Marginalia
+    module Comment
+      private
+
+      def jid
+        bg_job["jid"] if bg_job.present?
+      end
+
+      def job_class
+        bg_job["class"] if bg_job.present?
+      end
+
+      def correlation_id
+        if bg_job.present?
+          bg_job["correlation_id"]
+        else
+          Labkit::Correlation::CorrelationId.current_id
+        end
+      end
+
+      def bg_job
+        job = ::Marginalia::Comment.marginalia_job
+
+        # We are using 'Marginalia::SidekiqInstrumentation' which does not support 'ActiveJob::Base'.
+        # Gitlab also uses 'ActionMailer::DeliveryJob' which inherits from ActiveJob::Base.
+        # So below condition is used to return metadata for such jobs.
+        if job && job.is_a?(ActionMailer::DeliveryJob)
+          {
+            "class" => job.arguments.first,
+            "jid"   => job.job_id
+          }
+        else
+          job
+        end
+      end
+    end
+  end
+end