Inject correlation id

6 files changed, 67 insertions, 0 deletions

diff --git a/config/initializers/correlation_id.rb b/config/initializers/correlation_id.rb
new file mode 100644
index 00000000000..f1ebefbb0c5
--- /dev/null
+++ b/config/initializers/correlation_id.rb
@@ -0,0 +1 @@
+Rails.application.config.middleware.use(Gitlab::Middleware::CorrelationId)
diff --git a/config/initializers/sidekiq.rb b/config/initializers/sidekiq.rb
index 565efc858d1..e77e39c8485 100644
--- a/config/initializers/sidekiq.rb
+++ b/config/initializers/sidekiq.rb
@@ -20,6 +20,7 @@ Sidekiq.configure_server do |config|
     chain.add Gitlab::SidekiqMiddleware::ArgumentsLogger if ENV['SIDEKIQ_LOG_ARGUMENTS'] && !enable_json_logs
     chain.add Gitlab::SidekiqMiddleware::Shutdown
     chain.add Gitlab::SidekiqMiddleware::RequestStoreMiddleware unless ENV['SIDEKIQ_REQUEST_STORE'] == '0'
+    chain.add Gitlab::SidekiqMiddleware::CorrelationLogger
     chain.add Gitlab::SidekiqStatus::ServerMiddleware
   end
 
@@ -30,6 +31,7 @@ Sidekiq.configure_server do |config|
 
   config.client_middleware do |chain|
     chain.add Gitlab::SidekiqStatus::ClientMiddleware
+    chain.add Gitlab::SidekiqMiddleware::CorrelationInjector
   end
 
   config.on :startup do
diff --git a/lib/gitlab/json_logger.rb b/lib/gitlab/json_logger.rb
index 3bff77731f6..3564594e4da 100644
--- a/lib/gitlab/json_logger.rb
+++ b/lib/gitlab/json_logger.rb
@@ -6,6 +6,21 @@ module Gitlab
       raise NotImplementedError
     end
 
+    def self.use_correlation_id(correlation_id, &blk)
+      Thread.current[:correlation_id] ||= []
+      Thread.current[:correlation_id] << correlation_id
+
+      begin
+        yield
+      ensure
+        Thread.current[:correlation_id].unshift
+      end
+    end
+
+    def self.current_correlation_id
+      Thread.current[:correlation_id]&.last
+    end
+
     def format_message(severity, timestamp, progname, message)
       data = {}
       data[:severity] = severity
diff --git a/lib/gitlab/middleware/correlation_id.rb b/lib/gitlab/middleware/correlation_id.rb
new file mode 100644
index 00000000000..51e2a51ac1b
--- /dev/null
+++ b/lib/gitlab/middleware/correlation_id.rb
@@ -0,0 +1,27 @@
+# A dumb middleware that returns a Go HTML document if the go-get=1 query string
+# is used irrespective if the namespace/project exists
+module Gitlab
+  module Middleware
+    class CorrelationId
+      include ActionView::Helpers::TagHelper
+
+      CORRELATION_HEADER = 'HTTP_CORRELATION_ID'
+
+      def initialize(app)
+        @app = app
+      end
+
+      def call(env)
+        Gitlab::JsonLogger.use_correlation_id(correlation_id) do
+          @app.call(env)
+        end
+      end
+
+      private
+
+      def correlation_id
+        env[CORRELATION_HEADER] || SecureRandom.hex
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/sidekiq_middleware/correlation_injector.rb b/lib/gitlab/sidekiq_middleware/correlation_injector.rb
new file mode 100644
index 00000000000..4dc8646cbab
--- /dev/null
+++ b/lib/gitlab/sidekiq_middleware/correlation_injector.rb
@@ -0,0 +1,11 @@
+module Gitlab
+  module SidekiqMiddleware
+    class CorrelationInjector
+      def call(worker_class, job, queue, redis_pool)
+        job['correlation_id'] = Gitlab::JsonLogger.current_correlation_id
+
+        yield
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/sidekiq_middleware/correlation_logger.rb b/lib/gitlab/sidekiq_middleware/correlation_logger.rb
new file mode 100644
index 00000000000..3bd82449b2b
--- /dev/null
+++ b/lib/gitlab/sidekiq_middleware/correlation_logger.rb
@@ -0,0 +1,11 @@
+module Gitlab
+  module SidekiqMiddleware
+    class CorrelationLogger
+      def call(worker, job, queue)
+        Gitlab::JsonLogger.use_correlation_id("sidekiq:#{job['jid']}") do
+          yield
+        end
+      end
+    end
+  end
+end