Access feature category in the middleware by application context

Issue https://gitlab.com/gitlab-com/gl-infra/scalability/-/issues/1083

app/controllers/application_controller.rb

@@ -263,7 +263,6 @@ class ApplicationController < ActionController::Base
     headers['X-XSS-Protection'] = '1; mode=block'
     headers['X-UA-Compatible'] = 'IE=edge'
     headers['X-Content-Type-Options'] = 'nosniff'
-    headers[Gitlab::Metrics::RequestsRackMiddleware::FEATURE_CATEGORY_HEADER] = feature_category
   end
 
   def default_cache_headers
@@ -456,17 +455,17 @@ class ApplicationController < ActionController::Base
   end
 
   def set_current_context(&block)
-    Gitlab::ApplicationContext.with_context(
+    Gitlab::ApplicationContext.push(
       user: -> { context_user },
       project: -> { @project if @project&.persisted? },
       namespace: -> { @group if @group&.persisted? },
       caller_id: caller_id,
       remote_ip: request.ip,
-      feature_category: feature_category) do
-      yield
-    ensure
-      @current_context = Gitlab::ApplicationContext.current
-    end
+      feature_category: feature_category
+    )
+    yield
+  ensure
+    @current_context = Gitlab::ApplicationContext.current
   end
 
   def set_locale(&block)

lib/api/api.rb

@@ -52,8 +52,6 @@ module API
       api_endpoint = env['api.endpoint']
       feature_category = api_endpoint.options[:for].try(:feature_category_for_app, api_endpoint).to_s
 
-      header[Gitlab::Metrics::RequestsRackMiddleware::FEATURE_CATEGORY_HEADER] = feature_category
-
       Gitlab::ApplicationContext.push(
         user: -> { @current_user },
         project: -> { @project },

lib/api/helpers.rb

@@ -487,9 +487,8 @@ module API
     def handle_api_exception(exception)
       if report_exception?(exception)
         define_params_for_grape_middleware
-        Gitlab::ApplicationContext.with_context(user: current_user) do
-          Gitlab::ErrorTracking.track_exception(exception)
-        end
+        Gitlab::ApplicationContext.push(user: current_user)
+        Gitlab::ErrorTracking.track_exception(exception)
       end
 
       # This is used with GrapeLogging::Loggers::ExceptionLogger

lib/api/internal/base.rb

@@ -10,8 +10,6 @@ module API
         api_endpoint = env['api.endpoint']
         feature_category = api_endpoint.options[:for].try(:feature_category_for_app, api_endpoint).to_s
 
-        header[Gitlab::Metrics::RequestsRackMiddleware::FEATURE_CATEGORY_HEADER] = feature_category
-
         Gitlab::ApplicationContext.push(
           user: -> { actor&.user },
           project: -> { project },

lib/gitlab/application_context.rb

@@ -44,6 +44,10 @@ module Gitlab
       current.include?(Labkit::Context.log_key(attribute_name))
     end
 
+    def self.current_context_attribute(attribute_name)
+      Labkit::Context.current&.get_attribute(attribute_name)
+    end
+
     def initialize(**args)
       unknown_attributes = args.keys - APPLICATION_ATTRIBUTES.map(&:name)
       raise ArgumentError, "#{unknown_attributes} are not known keys" if unknown_attributes.any?

lib/gitlab/etag_caching/middleware.rb

@@ -67,10 +67,11 @@ module Gitlab
 
         add_instrument_for_cache_hit(status_code, route, request)
 
+        Gitlab::ApplicationContext.push(feature_category: route.feature_category)
+
         new_headers = {
           'ETag' => etag,
-          'X-Gitlab-From-Cache' => 'true',
-          ::Gitlab::Metrics::RequestsRackMiddleware::FEATURE_CATEGORY_HEADER => route.feature_category
+          'X-Gitlab-From-Cache' => 'true'
         }
 
         [status_code, new_headers, []]

lib/gitlab/metrics/requests_rack_middleware.rb

@@ -15,7 +15,6 @@ module Gitlab
 
       HEALTH_ENDPOINT = /^\/-\/(liveness|readiness|health|metrics)\/?$/.freeze
 
-      FEATURE_CATEGORY_HEADER = 'X-Gitlab-Feature-Category'
       FEATURE_CATEGORY_DEFAULT = 'unknown'
 
       # These were the top 5 categories at a point in time, chosen as a
@@ -70,13 +69,11 @@ module Gitlab
         started = Time.now.to_f
         health_endpoint = health_endpoint?(env['PATH_INFO'])
         status = 'undefined'
-        feature_category = nil
 
         begin
           status, headers, body = @app.call(env)
 
           elapsed = Time.now.to_f - started
-          feature_category = headers&.fetch(FEATURE_CATEGORY_HEADER, nil)
 
           unless health_endpoint
             RequestsRackMiddleware.http_request_duration_seconds.observe({ method: method }, elapsed)
@@ -104,6 +101,10 @@ module Gitlab
 
         HEALTH_ENDPOINT.match?(CGI.unescape(path))
       end
+
+      def feature_category
+        ::Gitlab::ApplicationContext.current_context_attribute(:feature_category)
+      end
     end
   end
 end

spec/lib/gitlab/application_context_spec.rb

@@ -68,6 +68,24 @@ RSpec.describe Gitlab::ApplicationContext do
     end
   end
 
+  describe '.current_context_attribute' do
+    it 'returns the raw attribute value' do
+      described_class.with_context(caller_id: "Hello") do
+        expect(described_class.current_context_attribute(:caller_id)).to be('Hello')
+      end
+    end
+
+    it 'returns the attribute value with meta prefix' do
+      described_class.with_context(feature_category: "secure") do
+        expect(described_class.current_context_attribute('meta.feature_category')).to be('secure')
+      end
+    end
+
+    it 'returns nil if the key was not present in the current context' do
+      expect(described_class.current_context_attribute(:caller_id)).to be(nil)
+    end
+  end
+
   describe '#to_lazy_hash' do
     let_it_be(:user) { create(:user) }
     let_it_be(:project) { create(:project) }

spec/lib/gitlab/etag_caching/middleware_spec.rb

@@ -33,7 +33,6 @@ RSpec.describe Gitlab::EtagCaching::Middleware, :clean_gitlab_redis_shared_state
 
       expect(headers['ETag']).to be_nil
       expect(headers['X-Gitlab-From-Cache']).to be_nil
-      expect(headers[::Gitlab::Metrics::RequestsRackMiddleware::FEATURE_CATEGORY_HEADER]).to be_nil
     end
 
     it 'passes status code from app' do
@@ -41,6 +40,12 @@ RSpec.describe Gitlab::EtagCaching::Middleware, :clean_gitlab_redis_shared_state
 
       expect(status).to eq app_status_code
     end
+
+    it 'does not set feature category attribute' do
+      expect(Gitlab::ApplicationContext).not_to receive(:push)
+
+      _, _, _ = middleware.call(build_request(path, if_none_match))
+    end
   end
 
   context 'when there is no ETag in store for given resource' do
@@ -164,8 +169,15 @@ RSpec.describe Gitlab::EtagCaching::Middleware, :clean_gitlab_redis_shared_state
     it 'sets correct headers' do
       _, headers, _ = middleware.call(build_request(path, if_none_match))
 
-      expect(headers).to include('X-Gitlab-From-Cache' => 'true',
-                                 ::Gitlab::Metrics::RequestsRackMiddleware::FEATURE_CATEGORY_HEADER => 'issue_tracking')
+      expect(headers).to include('X-Gitlab-From-Cache' => 'true')
+    end
+
+    it "pushes route's feature category to the context" do
+      expect(Gitlab::ApplicationContext).to receive(:push).with(
+        feature_category: 'issue_tracking'
+      )
+
+      _, _, _ = middleware.call(build_request(path, if_none_match))
     end
 
     it_behaves_like 'sends a process_action.action_controller notification', 304

spec/lib/gitlab/metrics/requests_rack_middleware_spec.rb

@@ -7,8 +7,16 @@ RSpec.describe Gitlab::Metrics::RequestsRackMiddleware, :aggregate_failures do
 
   subject { described_class.new(app) }
 
+  around do |example|
+    # Simulate application context middleware
+    # In fact, this middleware cleans up the contexts after a request lifecycle
+    ::Gitlab::ApplicationContext.with_context({}) do
+      example.run
+    end
+  end
+
   describe '#call' do
-    let(:status) { 100 }
+    let(:status) { 200 }
     let(:env) { { 'REQUEST_METHOD' => 'GET' } }
     let(:stack_result) { [status, {}, 'body'] }
 
@@ -91,9 +99,9 @@ RSpec.describe Gitlab::Metrics::RequestsRackMiddleware, :aggregate_failures do
     end
 
     context 'feature category header' do
-      context 'when a feature category header is present' do
+      context 'when a feature category context is present' do
         before do
-          allow(app).to receive(:call).and_return([200, { described_class::FEATURE_CATEGORY_HEADER => 'issue_tracking' }, nil])
+          ::Gitlab::ApplicationContext.push(feature_category: 'issue_tracking')
         end
 
         it 'adds the feature category to the labels for http_requests_total' do
@@ -113,11 +121,20 @@ RSpec.describe Gitlab::Metrics::RequestsRackMiddleware, :aggregate_failures do
         end
       end
 
-      context 'when the feature category header is an empty string' do
+      context 'when application raises an exception when the feature category context is present' do
         before do
-          allow(app).to receive(:call).and_return([200, { described_class::FEATURE_CATEGORY_HEADER => '' }, nil])
+          ::Gitlab::ApplicationContext.push(feature_category: 'issue_tracking')
+          allow(app).to receive(:call).and_raise(StandardError)
+        end
+
+        it 'adds the feature category to the labels for http_requests_total' do
+          expect(described_class).to receive_message_chain(:http_requests_total, :increment).with(method: 'get', status: 'undefined', feature_category: 'issue_tracking')
+
+          expect { subject.call(env) }.to raise_error(StandardError)
         end
+      end
 
+      context 'when the feature category context is not available' do
         it 'sets the feature category to unknown' do
           expect(described_class).to receive_message_chain(:http_requests_total, :increment).with(method: 'get', status: '200', feature_category: 'unknown')
           expect(described_class).not_to receive(:http_health_requests_total)