Merge branch '224630-forward-correlation-id-as-elasticsearch-opaque-id' into 'master'

Forward correlation ID to Elasticsearch as X-Opaque-Id header

See merge request gitlab-org/gitlab!35403

doc/development/elasticsearch.md

@@ -218,6 +218,16 @@ be used both locally in development and on any deployed GitLab instance to
 diagnose poor search performance. This will show the exact queries being made,
 which is useful to diagnose why a search might be slow.
 
+### Correlation ID and X-Opaque-Id
+
+Our [correlation
+ID](./distributed_tracing.md#developer-guidelines-for-working-with-correlation-ids)
+is forwarded by all requests from Rails to Elasticsearch as the
+[`X-Opaque-Id`](https://www.elastic.co/guide/en/elasticsearch/reference/current/tasks.html#_identifying_running_tasks)
+header which allows us to track any
+[tasks](https://www.elastic.co/guide/en/elasticsearch/reference/current/tasks.html)
+in the cluster back the request in GitLab.
+
 ## Troubleshooting
 
 ### Getting `flood stage disk watermark [95%] exceeded`

ee/spec/lib/gitlab/elastic/client_spec.rb

@@ -57,12 +57,15 @@ RSpec.describe Gitlab::Elastic::Client do
       end
 
       it 'signs_requests' do
+        # Mock the correlation ID (passed as header) to have deterministic signature
+        allow(Labkit::Correlation::CorrelationId).to receive(:current_or_new_id).and_return('new-correlation-id')
+
         stub_instance_credentials(creds_fail_response)
         travel_to(Time.parse('20170303T133952Z')) do
           stub_request(:get, 'http://example-elastic:9200/foo/_all/1')
             .with(
               headers: {
-                'Authorization'        => 'AWS4-HMAC-SHA256 Credential=0/20170303/us-east-1/es/aws4_request, SignedHeaders=content-type;host;x-amz-content-sha256;x-amz-date, Signature=4ba2aae19a476152dacf5a2191da67b0cf81b9d7152dab5c42b1bba701da19f1',
+                'Authorization'        => 'AWS4-HMAC-SHA256 Credential=0/20170303/us-east-1/es/aws4_request, SignedHeaders=content-type;host;x-amz-content-sha256;x-amz-date;x-opaque-id, Signature=c3180885fb19ca2cf4673a361aa47615dddd3ed52159fffcfeda9e732d7c91b8',
                 'Content-Type'         => 'application/json',
                 'X-Amz-Content-Sha256' => 'e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855',
                 'X-Amz-Date'           => '20170303T133952Z'

ee/spec/lib/gitlab/instrumentation/elasticsearch_transport_spec.rb

@@ -57,4 +57,30 @@ RSpec.describe ::Gitlab::Instrumentation::ElasticsearchTransportInterceptor, :el
     expect(::Gitlab::Instrumentation::ElasticsearchTransport.query_time).to be > 0
     expect(::Gitlab::Instrumentation::ElasticsearchTransport.detail_store).not_to be_empty
   end
+
+  it 'adds the labkit correlation id as X-Opaque-Id to all requests' do
+    allow(Labkit::Correlation::CorrelationId).to receive(:current_or_new_id).and_return('new-correlation-id')
+
+    elasticsearch_url = Gitlab::CurrentSettings.elasticsearch_url[0]
+    stub_request(:any, elasticsearch_url)
+
+    Project.__elasticsearch__.client
+      .perform_request(:get, '/')
+
+    expect(a_request(:get, /#{elasticsearch_url}/)
+      .with(headers: { 'X-Opaque-Id' => 'new-correlation-id' })).to have_been_made
+  end
+
+  it 'does not override the X-Opaque-Id if it is already present' do
+    allow(Labkit::Correlation::CorrelationId).to receive(:current_or_new_id).and_return('new-correlation-id')
+
+    elasticsearch_url = Gitlab::CurrentSettings.elasticsearch_url[0]
+    stub_request(:any, elasticsearch_url)
+
+    Project.__elasticsearch__.client
+      .perform_request(:get, '/', {}, nil, { 'X-Opaque-Id': 'original-opaque-id' })
+
+    expect(a_request(:get, /#{elasticsearch_url}/)
+      .with(headers: { 'X-Opaque-Id' => 'original-opaque-id' })).to have_been_made
+  end
 end

lib/gitlab/instrumentation/elasticsearch_transport.rb

@@ -5,8 +5,10 @@ require 'elasticsearch-transport'
 module Gitlab
   module Instrumentation
     module ElasticsearchTransportInterceptor
-      def perform_request(*args)
+      def perform_request(method, path, params = {}, body = nil, headers = nil)
         start = Time.now
+        headers = (headers || {})
+          .reverse_merge({ 'X-Opaque-Id': Labkit::Correlation::CorrelationId.current_or_new_id })
         super
       ensure
         if ::Gitlab::SafeRequestStore.active?
@@ -14,7 +16,7 @@ module Gitlab
 
           ::Gitlab::Instrumentation::ElasticsearchTransport.increment_request_count
           ::Gitlab::Instrumentation::ElasticsearchTransport.add_duration(duration)
-          ::Gitlab::Instrumentation::ElasticsearchTransport.add_call_details(duration, args)
+          ::Gitlab::Instrumentation::ElasticsearchTransport.add_call_details(duration, method, path, params, body)
         end
       end
     end
@@ -47,14 +49,14 @@ module Gitlab
         ::Gitlab::SafeRequestStore[ELASTICSEARCH_CALL_DURATION] += duration
       end
 
-      def self.add_call_details(duration, args)
+      def self.add_call_details(duration, method, path, params, body)
         return unless Gitlab::PerformanceBar.enabled_for_request?
 
         detail_store << {
-          method: args[0],
-          path: args[1],
-          params: args[2],
-          body: args[3],
+          method: method,
+          path: path,
+          params: params,
+          body: body,
           duration: duration,
           backtrace: ::Gitlab::BacktraceCleaner.clean_backtrace(caller)
         }