Don't try to enable the request store twice

When calling `#with_request_store` we enable the request store and
ensure it's disabled after the block completes.

However, since the request store is thread local, we can't enable and
disable it twice withing the same block. In regular circumstances,
this doesn't happen, since we only enable the request store in
middleware, and they both run in a different process:

- Using RequestStore::Middleware as a rack middleware for requests
- Using Gitlab::SidekiqMiddleware::RequestStoreMiddleware for
  sidekiq-server.

In specs this could happen when running Sidekiq::Testing.inline!. If
that was used within a spec that had the `:request_store` tag, or
within a feature/request spec that ran the rack middleware.

In this case, the Sidekiq middleware would disable and clear the
request store too soon, causing failures in the specs that counted on
the request store, for example specs counting queries, or checking
the number of Gitaly calls.

lib/gitlab/with_request_store.rb

@@ -2,12 +2,24 @@
 
 module Gitlab
   module WithRequestStore
-    def with_request_store
+    def with_request_store(&block)
+      # Skip enabling the request store if it was already active. Whatever
+      # instantiated the request store first is responsible for clearing it
+      return yield if RequestStore.active?
+
+      enabling_request_store(&block)
+    end
+
+    private
+
+    def enabling_request_store
       RequestStore.begin!
       yield
     ensure
       RequestStore.end!
       RequestStore.clear!
     end
+
+    extend self
   end
 end

spec/lib/gitlab/with_request_store_spec.rb

+# frozen_string_literal: true
+
+require 'fast_spec_helper'
+require 'request_store'
+
+describe Gitlab::WithRequestStore do
+  let(:fake_class) { Class.new { include Gitlab::WithRequestStore } }
+
+  subject(:object) { fake_class.new }
+
+  describe "#with_request_store" do
+    it 'starts a request store and yields control' do
+      expect(RequestStore).to receive(:begin!).ordered
+      expect(RequestStore).to receive(:end!).ordered
+      expect(RequestStore).to receive(:clear!).ordered
+
+      expect { |b| object.with_request_store(&b) }.to yield_control
+    end
+
+    it 'only starts a request store once when nested' do
+      expect(RequestStore).to receive(:begin!).ordered.once.and_call_original
+      expect(RequestStore).to receive(:end!).ordered.once.and_call_original
+      expect(RequestStore).to receive(:clear!).ordered.once.and_call_original
+
+      object.with_request_store do
+        expect { |b| object.with_request_store(&b) }.to yield_control
+      end
+    end
+  end
+end

spec/spec_helper.rb

@@ -286,12 +286,7 @@ RSpec.configure do |config|
   end
 
   config.around(:example, :request_store) do |example|
-    RequestStore.begin!
-
-    example.run
-
-    RequestStore.end!
-    RequestStore.clear!
+    Gitlab::WithRequestStore.with_request_store { example.run }
   end
 
   config.around do |example|