Revert "Merge branch 'zj-improve-gitaly-pb' into 'master'"

This reverts merge request !23140

changelogs/unreleased/zj-improve-gitaly-pb.yml

----
-title: Show what RPC is called in the performance bar
-merge_request: 23140
-author:
-type: other

lib/gitlab/git/blob.rb

 # frozen_string_literal: true
 
+# Gitaly note: JV: seems to be completely migrated (behind feature flags).
+
 module Gitlab
   module Git
     class Blob

lib/gitlab/git/repository.rb

@@ -885,6 +885,12 @@ module Gitlab
         Gitlab::GitalyClient::ConflictsService.new(self, our_commit_oid, their_commit_oid)
       end
 
+      def gitaly_migrate(method, status: Gitlab::GitalyClient::MigrationStatus::OPT_IN, &block)
+        wrapped_gitaly_errors do
+          Gitlab::GitalyClient.migrate(method, status: status, &block)
+        end
+      end
+
       def clean_stale_repository_files
         wrapped_gitaly_errors do
           gitaly_repository_client.cleanup if exists?

lib/gitlab/gitaly_client.rb

@@ -9,6 +9,11 @@ require 'grpc/health/v1/health_services_pb'
 module Gitlab
   module GitalyClient
     include Gitlab::Metrics::Methods
+    module MigrationStatus
+      DISABLED = 1
+      OPT_IN = 2
+      OPT_OUT = 3
+    end
 
     class TooManyInvocationsError < StandardError
       attr_reader :call_site, :invocation_count, :max_call_stack
@@ -26,7 +31,7 @@ module Gitlab
       end
     end
 
-    SERVER_VERSION_FILE = 'GITALY_SERVER_VERSION'
+    SERVER_VERSION_FILE = 'GITALY_SERVER_VERSION'.freeze
     MAXIMUM_GITALY_CALLS = 35
     CLIENT_NAME = (Sidekiq.server? ? 'gitlab-sidekiq' : 'gitlab-web').freeze
 
@@ -38,6 +43,11 @@ module Gitlab
 
     self.query_time = 0
 
+    define_histogram :gitaly_migrate_call_duration_seconds do
+      docstring "Gitaly migration call execution timings"
+      base_labels gitaly_enabled: nil, feature: nil
+    end
+
     define_histogram :gitaly_controller_action_duration_seconds do
       docstring "Gitaly endpoint histogram by controller and action combination"
       base_labels Gitlab::Metrics::Transaction::BASE_LABELS.merge(gitaly_service: nil, rpc: nil)
@@ -116,6 +126,7 @@ module Gitlab
     def self.call(storage, service, rpc, request, remote_storage: nil, timeout: nil)
       start = Gitlab::Metrics::System.monotonic_time
       request_hash = request.is_a?(Google::Protobuf::MessageExts) ? request.to_h : {}
+      @current_call_id ||= SecureRandom.uuid
 
       enforce_gitaly_request_limits(:call)
 
@@ -134,7 +145,9 @@ module Gitlab
         current_transaction_labels.merge(gitaly_service: service.to_s, rpc: rpc.to_s),
         duration)
 
-      add_call_details(feature: "#{service}##{rpc}", duration: duration, request: request_hash, rpc: rpc)
+      add_call_details(id: @current_call_id, feature: service, duration: duration, request: request_hash)
+
+      @current_call_id = nil
     end
 
     def self.handle_grpc_unavailable!(ex)
@@ -209,7 +222,7 @@ module Gitlab
       result
     end
 
-    SERVER_FEATURE_FLAGS = %w[].freeze
+    SERVER_FEATURE_FLAGS = %w[gogit_findcommit].freeze
 
     def self.server_feature_flags
       SERVER_FEATURE_FLAGS.map do |f|
@@ -224,8 +237,82 @@ module Gitlab
       params['gitaly_token'].presence || Gitlab.config.gitaly['token']
     end
 
-    def self.feature_enabled?(feature_name)
-      Feature.enabled?("gitaly_#{feature_name}")
+    # Evaluates whether a feature toggle is on or off
+    def self.feature_enabled?(feature_name, status: MigrationStatus::OPT_IN)
+      # Disabled features are always off!
+      return false if status == MigrationStatus::DISABLED
+
+      feature = Feature.get("gitaly_#{feature_name}")
+
+      # If the feature has been set, always evaluate
+      if Feature.persisted?(feature)
+        if feature.percentage_of_time_value > 0
+          # Probabilistically enable this feature
+          return Random.rand() * 100 < feature.percentage_of_time_value
+        end
+
+        return feature.enabled?
+      end
+
+      # If the feature has not been set, the default depends
+      # on it's status
+      case status
+      when MigrationStatus::OPT_OUT
+        true
+      when MigrationStatus::OPT_IN
+        opt_into_all_features? && !explicit_opt_in_required.include?(feature_name)
+      else
+        false
+      end
+    rescue => ex
+      # During application startup feature lookups in SQL can fail
+      Rails.logger.warn "exception while checking Gitaly feature status for #{feature_name}: #{ex}"
+      false
+    end
+
+    # We have a mechanism to let GitLab automatically opt in to all Gitaly
+    # features. We want to be able to exclude some features from automatic
+    # opt-in. This function has an override in EE.
+    def self.explicit_opt_in_required
+      []
+    end
+
+    # opt_into_all_features? returns true when the current environment
+    # is one in which we opt into features automatically
+    def self.opt_into_all_features?
+      Rails.env.development? || ENV["GITALY_FEATURE_DEFAULT_ON"] == "1"
+    end
+    private_class_method :opt_into_all_features?
+
+    def self.migrate(feature, status: MigrationStatus::OPT_IN)
+      # Enforce limits at both the `migrate` and `call` sites to ensure that
+      # problems are not hidden by a feature being disabled
+      enforce_gitaly_request_limits(:migrate)
+
+      is_enabled  = feature_enabled?(feature, status: status)
+      metric_name = feature.to_s
+      metric_name += "_gitaly" if is_enabled
+
+      Gitlab::Metrics.measure(metric_name) do
+        # Some migrate calls wrap other migrate calls
+        allow_n_plus_1_calls do
+          feature_stack = Thread.current[:gitaly_feature_stack] ||= []
+          feature_stack.unshift(feature)
+          begin
+            start = Gitlab::Metrics::System.monotonic_time
+            @current_call_id = SecureRandom.uuid
+            call_details = { id: @current_call_id }
+            yield is_enabled
+          ensure
+            total_time = Gitlab::Metrics::System.monotonic_time - start
+            gitaly_migrate_call_duration_seconds.observe({ gitaly_enabled: is_enabled, feature: feature }, total_time)
+            feature_stack.shift
+            Thread.current[:gitaly_feature_stack] = nil if feature_stack.empty?
+
+            add_call_details(call_details.merge(feature: feature, duration: total_time))
+          end
+        end
+      end
     end
 
     # Ensures that Gitaly is not being abuse through n+1 misuse etc
@@ -281,20 +368,38 @@ module Gitlab
     end
     private_class_method :decrement_call_count
 
-    # Returns the of the number of Gitaly calls made for this request
+    # Returns an estimate of the number of Gitaly calls made for this
+    # request
     def self.get_request_count
-      get_call_count("gitaly_call_actual")
+      return 0 unless Gitlab::SafeRequestStore.active?
+
+      gitaly_migrate_count = get_call_count("gitaly_migrate_actual")
+      gitaly_call_count = get_call_count("gitaly_call_actual")
+
+      # Using the maximum of migrate and call_count will provide an
+      # indicator of how many Gitaly calls will be made, even
+      # before a feature is enabled. This provides us with a single
+      # metric, but not an exact number, but this tradeoff is acceptable
+      if gitaly_migrate_count > gitaly_call_count
+        gitaly_migrate_count
+      else
+        gitaly_call_count
+      end
     end
 
     def self.reset_counts
       return unless Gitlab::SafeRequestStore.active?
 
-      Gitlab::SafeRequestStore["gitaly_call_actual"] = 0
-      Gitlab::SafeRequestStore["gitaly_call_permitted"] = 0
+      %w[migrate call].each do |call_site|
+        Gitlab::SafeRequestStore["gitaly_#{call_site}_actual"] = 0
+        Gitlab::SafeRequestStore["gitaly_#{call_site}_permitted"] = 0
+      end
     end
 
     def self.add_call_details(details)
-      return unless Gitlab::SafeRequestStore[:peek_enabled]
+      id = details.delete(:id)
+
+      return unless id && Gitlab::SafeRequestStore[:peek_enabled]
 
       Gitlab::SafeRequestStore['gitaly_call_details'] ||= {}
       Gitlab::SafeRequestStore['gitaly_call_details'][id] ||= {}

spec/features/projects/wiki/user_views_wiki_page_spec.rb

 require 'spec_helper'
 
 describe 'User views a wiki page' do
-  include WikiHelpers
-
-  let(:user) { create(:user) }
-  let(:project) { create(:project, :wiki_repo, namespace: user.namespace) }
-  let(:path) { 'image.png' }
-  let(:wiki_page) do
-    create(:wiki_page,
-           wiki: project.wiki,
-           attrs: { title: 'home', content: "Look at this [image](#{path})\n\n ![alt text](#{path})" })
-  end
-
-  before do
-    project.add_maintainer(user)
-    sign_in(user)
-  end
+  shared_examples 'wiki page user view' do
+    include WikiHelpers
+
+    let(:user) { create(:user) }
+    let(:project) { create(:project, :wiki_repo, namespace: user.namespace) }
+    let(:path) { 'image.png' }
+    let(:wiki_page) do
+      create(:wiki_page,
+        wiki: project.wiki,
+        attrs: { title: 'home', content: "Look at this [image](#{path})\n\n ![alt text](#{path})" })
+    end
 
-  context 'when wiki is empty' do
     before do
-      visit(project_wikis_path(project))
-      click_link "Create your first page"
+      project.add_maintainer(user)
+      sign_in(user)
+    end
 
-      click_on('New page')
+    context 'when wiki is empty' do
+      before do
+        visit(project_wikis_path(project))
+        click_link "Create your first page"
 
-      page.within('#modal-new-wiki') do
-        fill_in(:new_wiki_path, with: 'one/two/three-test')
-        click_on('Create page')
-      end
+        click_on('New page')
+
+        page.within('#modal-new-wiki') do
+          fill_in(:new_wiki_path, with: 'one/two/three-test')
+          click_on('Create page')
+        end
 
-      page.within('.wiki-form') do
-        fill_in(:wiki_content, with: 'wiki content')
-        click_on('Create page')
+        page.within('.wiki-form') do
+          fill_in(:wiki_content, with: 'wiki content')
+          click_on('Create page')
+        end
       end
-    end
 
-    it 'shows the history of a page that has a path', :js do
-      expect(current_path).to include('one/two/three-test')
+      it 'shows the history of a page that has a path', :js do
+        expect(current_path).to include('one/two/three-test')
 
-      first(:link, text: 'Three').click
-      click_on('Page history')
+        first(:link, text: 'Three').click
+        click_on('Page history')
 
-      expect(current_path).to include('one/two/three-test')
+        expect(current_path).to include('one/two/three-test')
 
-      page.within(:css, '.nav-text') do
-        expect(page).to have_content('History')
+        page.within(:css, '.nav-text') do
+          expect(page).to have_content('History')
+        end
       end
-    end
 
-    it 'shows an old version of a page', :js do
-      expect(current_path).to include('one/two/three-test')
-      expect(find('.wiki-pages')).to have_content('Three')
+      it 'shows an old version of a page', :js do
+        expect(current_path).to include('one/two/three-test')
+        expect(find('.wiki-pages')).to have_content('Three')
 
-      first(:link, text: 'Three').click
+        first(:link, text: 'Three').click
 
-      expect(find('.nav-text')).to have_content('Three')
+        expect(find('.nav-text')).to have_content('Three')
 
-      click_on('Edit')
+        click_on('Edit')
 
-      expect(current_path).to include('one/two/three-test')
-      expect(page).to have_content('Edit Page')
+        expect(current_path).to include('one/two/three-test')
+        expect(page).to have_content('Edit Page')
 
-      fill_in('Content', with: 'Updated Wiki Content')
+        fill_in('Content', with: 'Updated Wiki Content')
 
-      click_on('Save changes')
-      click_on('Page history')
+        click_on('Save changes')
+        click_on('Page history')
 
-      page.within(:css, '.nav-text') do
-        expect(page).to have_content('History')
-      end
+        page.within(:css, '.nav-text') do
+          expect(page).to have_content('History')
+        end
 
-      find('a[href*="?version_id"]')
+        find('a[href*="?version_id"]')
+      end
     end
-  end
 
-  context 'when a page does not have history' do
-    before do
-      visit(project_wiki_path(project, wiki_page))
-    end
+    context 'when a page does not have history' do
+      before do
+        visit(project_wiki_path(project, wiki_page))
+      end
 
-    it 'shows all the pages' do
-      expect(page).to have_content(user.name)
-      expect(find('.wiki-pages')).to have_content(wiki_page.title.capitalize)
-    end
+      it 'shows all the pages' do
+        expect(page).to have_content(user.name)
+        expect(find('.wiki-pages')).to have_content(wiki_page.title.capitalize)
+      end
+
+      context 'shows a file stored in a page' do
+        let(:path) { upload_file_to_wiki(project, user, 'dk.png') }
+
+        it do
+          expect(page).to have_xpath("//img[@data-src='#{project.wiki.wiki_base_path}/#{path}']")
+          expect(page).to have_link('image', href: "#{project.wiki.wiki_base_path}/#{path}")
+
+          click_on('image')
 
-    context 'shows a file stored in a page' do
-      let(:path) { upload_file_to_wiki(project, user, 'dk.png') }
+          expect(current_path).to match("wikis/#{path}")
+          expect(page).not_to have_xpath('/html') # Page should render the image which means there is no html involved
+        end
+      end
 
-      it do
-        expect(page).to have_xpath("//img[@data-src='#{project.wiki.wiki_base_path}/#{path}']")
+      it 'shows the creation page if file does not exist' do
         expect(page).to have_link('image', href: "#{project.wiki.wiki_base_path}/#{path}")
 
         click_on('image')
 
         expect(current_path).to match("wikis/#{path}")
-        expect(page).not_to have_xpath('/html') # Page should render the image which means there is no html involved
+        expect(page).to have_content('New Wiki Page')
+        expect(page).to have_content('Create page')
       end
     end
 
-    it 'shows the creation page if file does not exist' do
-      expect(page).to have_link('image', href: "#{project.wiki.wiki_base_path}/#{path}")
+    context 'when a page has history' do
+      before do
+        wiki_page.update(message: 'updated home', content: 'updated [some link](other-page)')
+      end
 
-      click_on('image')
+      it 'shows the page history' do
+        visit(project_wiki_path(project, wiki_page))
 
-      expect(current_path).to match("wikis/#{path}")
-      expect(page).to have_content('New Wiki Page')
-      expect(page).to have_content('Create page')
-    end
-  end
+        expect(page).to have_selector('a.btn', text: 'Edit')
 
-  context 'when a page has history' do
-    before do
-      wiki_page.update(message: 'updated home', content: 'updated [some link](other-page)')
-    end
+        click_on('Page history')
 
-    it 'shows the page history' do
-      visit(project_wiki_path(project, wiki_page))
-
-      expect(page).to have_selector('a.btn', text: 'Edit')
+        expect(page).to have_content(user.name)
+        expect(page).to have_content("#{user.username} created page: home")
+        expect(page).to have_content('updated home')
+      end
 
-      click_on('Page history')
+      it 'does not show the "Edit" button' do
+        visit(project_wiki_path(project, wiki_page, version_id: wiki_page.versions.last.id))
 
-      expect(page).to have_content(user.name)
-      expect(page).to have_content("#{user.username} created page: home")
-      expect(page).to have_content('updated home')
+        expect(page).not_to have_selector('a.btn', text: 'Edit')
+      end
     end
 
-    it 'does not show the "Edit" button' do
-      visit(project_wiki_path(project, wiki_page, version_id: wiki_page.versions.last.id))
+    context 'when page has invalid content encoding' do
+      let(:content) { 'whatever'.force_encoding('ISO-8859-1') }
 
-      expect(page).not_to have_selector('a.btn', text: 'Edit')
-    end
-  end
+      before do
+        allow(Gitlab::EncodingHelper).to receive(:encode!).and_return(content)
 
-  context 'when page has invalid content encoding' do
-    let(:content) { 'whatever'.force_encoding('ISO-8859-1') }
+        visit(project_wiki_path(project, wiki_page))
+      end
 
-    before do
-      allow(Gitlab::EncodingHelper).to receive(:encode!).and_return(content)
+      it 'does not show "Edit" button' do
+        expect(page).not_to have_selector('a.btn', text: 'Edit')
+      end
 
-      visit(project_wiki_path(project, wiki_page))
+      it 'shows error' do
+        page.within(:css, '.flash-notice') do
+          expect(page).to have_content('The content of this page is not encoded in UTF-8. Edits can only be made via the Git repository.')
+        end
+      end
     end
 
-    it 'does not show "Edit" button' do
-      expect(page).not_to have_selector('a.btn', text: 'Edit')
-    end
+    it 'opens a default wiki page', :js do
+      visit(project_path(project))
 
-    it 'shows error' do
-      page.within(:css, '.flash-notice') do
-        expect(page).to have_content('The content of this page is not encoded in UTF-8. Edits can only be made via the Git repository.')
-      end
+      find('.shortcuts-wiki').click
+      click_link "Create your first page"
+
+      expect(page).to have_content('Home · Create Page')
     end
   end
 
-  it 'opens a default wiki page', :js do
-    visit(project_path(project))
-
-    find('.shortcuts-wiki').click
-    click_link "Create your first page"
+  context 'when Gitaly is enabled' do
+    it_behaves_like 'wiki page user view'
+  end
 
-    expect(page).to have_content('Home · Create Page')
+  context 'when Gitaly is disabled', :skip_gitaly_mock do
+    it_behaves_like 'wiki page user view'
   end
 end

spec/lib/extracts_path_spec.rb

@@ -205,18 +205,28 @@ describe ExtractsPath do
   end
 
   describe '#lfs_blob_ids' do
-    let(:tag) { @project.repository.add_tag(@project.owner, 'my-annotated-tag', 'master', 'test tag') }
-    let(:ref) { tag.target }
-    let(:params) { { ref: ref, path: 'README.md' } }
+    shared_examples '#lfs_blob_ids' do
+      let(:tag) { @project.repository.add_tag(@project.owner, 'my-annotated-tag', 'master', 'test tag') }
+      let(:ref) { tag.target }
+      let(:params) { { ref: ref, path: 'README.md' } }
 
-    before do
-      @project = create(:project, :repository)
+      before do
+        @project = create(:project, :repository)
+      end
+
+      it 'handles annotated tags' do
+        assign_ref_vars
+
+        expect(lfs_blob_ids).to eq([])
+      end
     end
 
-    it 'handles annotated tags' do
-      assign_ref_vars
+    context 'when gitaly is enabled' do
+      it_behaves_like '#lfs_blob_ids'
+    end
 
-      expect(lfs_blob_ids).to eq([])
+    context 'when gitaly is disabled', :skip_gitaly_mock do
+      it_behaves_like '#lfs_blob_ids'
     end
   end
 end

spec/lib/gitlab/diff/file_collection/merge_request_diff_spec.rb

@@ -37,7 +37,17 @@ describe Gitlab::Diff::FileCollection::MergeRequestDiff do
     let(:stub_path) { '.gitignore' }
   end
 
-  it 'returns a valid instance of a DiffCollection' do
-    expect(diff_files).to be_a(Gitlab::Git::DiffCollection)
+  shared_examples 'initializes a DiffCollection' do
+    it 'returns a valid instance of a DiffCollection' do
+      expect(diff_files).to be_a(Gitlab::Git::DiffCollection)
+    end
+  end
+
+  context 'with Gitaly disabled', :disable_gitaly do
+    it_behaves_like 'initializes a DiffCollection'
+  end
+
+  context 'with Gitaly enabled' do
+    it_behaves_like 'initializes a DiffCollection'
   end
 end

spec/lib/gitlab/git/blob_spec.rb

@@ -128,7 +128,7 @@ describe Gitlab::Git::Blob, :seed_helper do
     end
   end
 
-  describe '.raw' do
+  shared_examples 'finding blobs by ID' do
     let(:raw_blob) { Gitlab::Git::Blob.raw(repository, SeedRepo::RubyBlob::ID) }
     let(:bad_blob) { Gitlab::Git::Blob.raw(repository, SeedRepo::BigCommit::ID) }
 
@@ -166,6 +166,16 @@ describe Gitlab::Git::Blob, :seed_helper do
     end
   end
 
+  describe '.raw' do
+    context 'when the blob_raw Gitaly feature is enabled' do
+      it_behaves_like 'finding blobs by ID'
+    end
+
+    context 'when the blob_raw Gitaly feature is disabled', :skip_gitaly_mock do
+      it_behaves_like 'finding blobs by ID'
+    end
+  end
+
   describe '.batch' do
     let(:blob_references) do
       [

spec/lib/gitlab/git/commit_spec.rb

@@ -183,100 +183,110 @@ describe Gitlab::Git::Commit, :seed_helper do
       end
     end
 
-    context 'path is empty string' do
-      subject do
-        commits = described_class.where(
-          repo: repository,
-          ref: 'master',
-          path: '',
-          limit: 10
-        )
+    shared_examples '.where' do
+      context 'path is empty string' do
+        subject do
+          commits = described_class.where(
+            repo: repository,
+            ref: 'master',
+            path: '',
+            limit: 10
+          )
 
-        commits.map { |c| c.id }
-      end
+          commits.map { |c| c.id }
+        end
 
-      it 'has 10 elements' do
-        expect(subject.size).to eq(10)
+        it 'has 10 elements' do
+          expect(subject.size).to eq(10)
+        end
+        it { is_expected.to include(SeedRepo::EmptyCommit::ID) }
       end
-      it { is_expected.to include(SeedRepo::EmptyCommit::ID) }
-    end
 
-    context 'path is nil' do
-      subject do
-        commits = described_class.where(
-          repo: repository,
-          ref: 'master',
-          path: nil,
-          limit: 10
-        )
+      context 'path is nil' do
+        subject do
+          commits = described_class.where(
+            repo: repository,
+            ref: 'master',
+            path: nil,
+            limit: 10
+          )
 
-        commits.map { |c| c.id }
-      end
+          commits.map { |c| c.id }
+        end
 
-      it 'has 10 elements' do
-        expect(subject.size).to eq(10)
+        it 'has 10 elements' do
+          expect(subject.size).to eq(10)
+        end
+        it { is_expected.to include(SeedRepo::EmptyCommit::ID) }
       end
-      it { is_expected.to include(SeedRepo::EmptyCommit::ID) }
-    end
 
-    context 'ref is branch name' do
-      subject do
-        commits = described_class.where(
-          repo: repository,
-          ref: 'master',
-          path: 'files',
-          limit: 3,
-          offset: 1
-        )
+      context 'ref is branch name' do
+        subject do
+          commits = described_class.where(
+            repo: repository,
+            ref: 'master',
+            path: 'files',
+            limit: 3,
+            offset: 1
+          )
 
-        commits.map { |c| c.id }
-      end
+          commits.map { |c| c.id }
+        end
 
-      it 'has 3 elements' do
-        expect(subject.size).to eq(3)
+        it 'has 3 elements' do
+          expect(subject.size).to eq(3)
+        end
+        it { is_expected.to include("d14d6c0abdd253381df51a723d58691b2ee1ab08") }
+        it { is_expected.not_to include("eb49186cfa5c4338011f5f590fac11bd66c5c631") }
       end
-      it { is_expected.to include("d14d6c0abdd253381df51a723d58691b2ee1ab08") }
-      it { is_expected.not_to include("eb49186cfa5c4338011f5f590fac11bd66c5c631") }
-    end
 
-    context 'ref is commit id' do
-      subject do
-        commits = described_class.where(
-          repo: repository,
-          ref: "874797c3a73b60d2187ed6e2fcabd289ff75171e",
-          path: 'files',
-          limit: 3,
-          offset: 1
-        )
+      context 'ref is commit id' do
+        subject do
+          commits = described_class.where(
+            repo: repository,
+            ref: "874797c3a73b60d2187ed6e2fcabd289ff75171e",
+            path: 'files',
+            limit: 3,
+            offset: 1
+          )
 
-        commits.map { |c| c.id }
-      end
+          commits.map { |c| c.id }
+        end
 
-      it 'has 3 elements' do
-        expect(subject.size).to eq(3)
+        it 'has 3 elements' do
+          expect(subject.size).to eq(3)
+        end
+        it { is_expected.to include("2f63565e7aac07bcdadb654e253078b727143ec4") }
+        it { is_expected.not_to include(SeedRepo::Commit::ID) }
       end
-      it { is_expected.to include("2f63565e7aac07bcdadb654e253078b727143ec4") }
-      it { is_expected.not_to include(SeedRepo::Commit::ID) }
-    end
 
-    context 'ref is tag' do
-      subject do
-        commits = described_class.where(
-          repo: repository,
-          ref: 'v1.0.0',
-          path: 'files',
-          limit: 3,
-          offset: 1
-        )
+      context 'ref is tag' do
+        subject do
+          commits = described_class.where(
+            repo: repository,
+            ref: 'v1.0.0',
+            path: 'files',
+            limit: 3,
+            offset: 1
+          )
 
-        commits.map { |c| c.id }
-      end
+          commits.map { |c| c.id }
+        end
 
-      it 'has 3 elements' do
-        expect(subject.size).to eq(3)
+        it 'has 3 elements' do
+          expect(subject.size).to eq(3)
+        end
+        it { is_expected.to include("874797c3a73b60d2187ed6e2fcabd289ff75171e") }
+        it { is_expected.not_to include(SeedRepo::Commit::ID) }
       end
-      it { is_expected.to include("874797c3a73b60d2187ed6e2fcabd289ff75171e") }
-      it { is_expected.not_to include(SeedRepo::Commit::ID) }
+    end
+
+    describe '.where with gitaly' do
+      it_should_behave_like '.where'
+    end
+
+    describe '.where without gitaly', :skip_gitaly_mock do
+      it_should_behave_like '.where'
     end
 
     describe '.between' do
@@ -498,7 +508,7 @@ describe Gitlab::Git::Commit, :seed_helper do
     end
   end
 
-  describe '#stats' do
+  shared_examples '#stats' do
     subject { commit.stats }
 
     describe '#additions' do
@@ -517,6 +527,14 @@ describe Gitlab::Git::Commit, :seed_helper do
     end
   end
 
+  describe '#stats with gitaly on' do
+    it_should_behave_like '#stats'
+  end
+
+  describe '#stats with gitaly disabled', :skip_gitaly_mock do
+    it_should_behave_like '#stats'
+  end
+
   describe '#has_zero_stats?' do
     it { expect(commit.has_zero_stats?).to eq(false) }
   end
@@ -559,15 +577,25 @@ describe Gitlab::Git::Commit, :seed_helper do
       commit_ids.map { |id| described_class.get_message(repository, id) }
     end
 
-    it 'gets commit messages' do
-      expect(subject).to contain_exactly(
-        "Added contributing guide\n\nSigned-off-by: Dmitriy Zaporozhets <dmitriy.zaporozhets@gmail.com>\n",
-        "Add submodule\n\nSigned-off-by: Dmitriy Zaporozhets <dmitriy.zaporozhets@gmail.com>\n"
-      )
+    shared_examples 'getting commit messages' do
+      it 'gets commit messages' do
+        expect(subject).to contain_exactly(
+          "Added contributing guide\n\nSigned-off-by: Dmitriy Zaporozhets <dmitriy.zaporozhets@gmail.com>\n",
+          "Add submodule\n\nSigned-off-by: Dmitriy Zaporozhets <dmitriy.zaporozhets@gmail.com>\n"
+        )
+      end
+    end
+
+    context 'when Gitaly commit_messages feature is enabled' do
+      it_behaves_like 'getting commit messages'
+
+      it 'gets messages in one batch', :request_store do
+        expect { subject.map(&:itself) }.to change { Gitlab::GitalyClient.get_request_count }.by(1)
+      end
     end
 
-    it 'gets messages in one batch', :request_store do
-      expect { subject.map(&:itself) }.to change { Gitlab::GitalyClient.get_request_count }.by(1)
+    context 'when Gitaly commit_messages feature is disabled', :disable_gitaly do
+      it_behaves_like 'getting commit messages'
     end
   end
 

spec/lib/gitlab/git/tag_spec.rb

@@ -3,7 +3,7 @@ require "spec_helper"
 describe Gitlab::Git::Tag, :seed_helper do
   let(:repository) { Gitlab::Git::Repository.new('default', TEST_REPO_PATH, '') }
 
-  describe '#tags' do
+  shared_examples 'Gitlab::Git::Repository#tags' do
     describe 'first tag' do
       let(:tag) { repository.tags.first }
 
@@ -25,6 +25,14 @@ describe Gitlab::Git::Tag, :seed_helper do
     it { expect(repository.tags.size).to eq(SeedRepo::Repo::TAGS.size) }
   end
 
+  context 'when Gitaly tags feature is enabled' do
+    it_behaves_like 'Gitlab::Git::Repository#tags'
+  end
+
+  context 'when Gitaly tags feature is disabled', :skip_gitaly_mock do
+    it_behaves_like 'Gitlab::Git::Repository#tags'
+  end
+
   describe '.get_message' do
     let(:tag_ids) { %w[f4e6814c3e4e7a0de82a9e7cd20c626cc963a2f8 8a2a6eb295bb170b34c24c76c49ed0e9b2eaf34b] }
 
@@ -32,13 +40,23 @@ describe Gitlab::Git::Tag, :seed_helper do
       tag_ids.map { |id| described_class.get_message(repository, id) }
     end
 
-    it 'gets tag messages' do
-      expect(subject[0]).to eq("Release\n")
-      expect(subject[1]).to eq("Version 1.1.0\n")
+    shared_examples 'getting tag messages' do
+      it 'gets tag messages' do
+        expect(subject[0]).to eq("Release\n")
+        expect(subject[1]).to eq("Version 1.1.0\n")
+      end
+    end
+
+    context 'when Gitaly tag_messages feature is enabled' do
+      it_behaves_like 'getting tag messages'
+
+      it 'gets messages in one batch', :request_store do
+        expect { subject.map(&:itself) }.to change { Gitlab::GitalyClient.get_request_count }.by(1)
+      end
     end
 
-    it 'gets messages in one batch', :request_store do
-      expect { subject.map(&:itself) }.to change { Gitlab::GitalyClient.get_request_count }.by(1)
+    context 'when Gitaly tag_messages feature is disabled', :disable_gitaly do
+      it_behaves_like 'getting tag messages'
     end
   end
 

spec/lib/gitlab/git/tree_spec.rb

@@ -80,8 +80,18 @@ describe Gitlab::Git::Tree, :seed_helper do
   end
 
   describe '#where' do
-    it 'returns an empty array when called with an invalid ref' do
-      expect(described_class.where(repository, 'foobar-does-not-exist')).to eq([])
+    shared_examples '#where' do
+      it 'returns an empty array when called with an invalid ref' do
+        expect(described_class.where(repository, 'foobar-does-not-exist')).to eq([])
+      end
+    end
+
+    context 'with gitaly' do
+      it_behaves_like '#where'
+    end
+
+    context 'without gitaly', :skip_gitaly_mock do
+      it_behaves_like '#where'
     end
   end
 end

spec/lib/gitlab/gitaly_client_spec.rb

@@ -2,7 +2,7 @@ require 'spec_helper'
 
 # We stub Gitaly in `spec/support/gitaly.rb` for other tests. We don't want
 # those stubs while testing the GitalyClient itself.
-describe Gitlab::GitalyClient do
+describe Gitlab::GitalyClient, skip_gitaly_mock: true do
   describe '.stub_class' do
     it 'returns the gRPC health check stub' do
       expect(described_class.stub_class(:health_check)).to eq(::Grpc::Health::V1::Health::Stub)
@@ -191,13 +191,102 @@ describe Gitlab::GitalyClient do
     let(:feature_name) { 'my_feature' }
     let(:real_feature_name) { "gitaly_#{feature_name}" }
 
-    before do
-      allow(Feature).to receive(:enabled?).and_return(false)
+    context 'when Gitaly is disabled' do
+      before do
+        allow(described_class).to receive(:enabled?).and_return(false)
+      end
+
+      it 'returns false' do
+        expect(described_class.feature_enabled?(feature_name)).to be(false)
+      end
     end
 
-    it 'returns false' do
-      expect(Feature).to receive(:enabled?).with(real_feature_name)
-      expect(described_class.feature_enabled?(feature_name)).to be(false)
+    context 'when the feature status is DISABLED' do
+      let(:feature_status) { Gitlab::GitalyClient::MigrationStatus::DISABLED }
+
+      it 'returns false' do
+        expect(described_class.feature_enabled?(feature_name, status: feature_status)).to be(false)
+      end
+    end
+
+    context 'when the feature_status is OPT_IN' do
+      let(:feature_status) { Gitlab::GitalyClient::MigrationStatus::OPT_IN }
+
+      context "when the feature flag hasn't been set" do
+        it 'returns false' do
+          expect(described_class.feature_enabled?(feature_name, status: feature_status)).to be(false)
+        end
+      end
+
+      context "when the feature flag is set to disable" do
+        before do
+          Feature.get(real_feature_name).disable
+        end
+
+        it 'returns false' do
+          expect(described_class.feature_enabled?(feature_name, status: feature_status)).to be(false)
+        end
+      end
+
+      context "when the feature flag is set to enable" do
+        before do
+          Feature.get(real_feature_name).enable
+        end
+
+        it 'returns true' do
+          expect(described_class.feature_enabled?(feature_name, status: feature_status)).to be(true)
+        end
+      end
+
+      context "when the feature flag is set to a percentage of time" do
+        before do
+          Feature.get(real_feature_name).enable_percentage_of_time(70)
+        end
+
+        it 'bases the result on pseudo-random numbers' do
+          expect(Random).to receive(:rand).and_return(0.3)
+          expect(described_class.feature_enabled?(feature_name, status: feature_status)).to be(true)
+
+          expect(Random).to receive(:rand).and_return(0.8)
+          expect(described_class.feature_enabled?(feature_name, status: feature_status)).to be(false)
+        end
+      end
+
+      context "when a feature is not persisted" do
+        it 'returns false when opt_into_all_features is off' do
+          allow(Feature).to receive(:persisted?).and_return(false)
+          allow(described_class).to receive(:opt_into_all_features?).and_return(false)
+
+          expect(described_class.feature_enabled?(feature_name, status: feature_status)).to be(false)
+        end
+
+        it 'returns true when the override is on' do
+          allow(Feature).to receive(:persisted?).and_return(false)
+          allow(described_class).to receive(:opt_into_all_features?).and_return(true)
+
+          expect(described_class.feature_enabled?(feature_name, status: feature_status)).to be(true)
+        end
+      end
+    end
+
+    context 'when the feature_status is OPT_OUT' do
+      let(:feature_status) { Gitlab::GitalyClient::MigrationStatus::OPT_OUT }
+
+      context "when the feature flag hasn't been set" do
+        it 'returns true' do
+          expect(described_class.feature_enabled?(feature_name, status: feature_status)).to be(true)
+        end
+      end
+
+      context "when the feature flag is set to disable" do
+        before do
+          Feature.get(real_feature_name).disable
+        end
+
+        it 'returns false' do
+          expect(described_class.feature_enabled?(feature_name, status: feature_status)).to be(false)
+        end
+      end
     end
   end
 

spec/models/merge_request_spec.rb

@@ -30,38 +30,48 @@ describe MergeRequest do
   end
 
   describe '#squash_in_progress?' do
-    let(:repo_path) do
-      Gitlab::GitalyClient::StorageSettings.allow_disk_access do
-        subject.source_project.repository.path
+    shared_examples 'checking whether a squash is in progress' do
+      let(:repo_path) do
+        Gitlab::GitalyClient::StorageSettings.allow_disk_access do
+          subject.source_project.repository.path
+        end
       end
-    end
-    let(:squash_path) { File.join(repo_path, "gitlab-worktree", "squash-#{subject.id}") }
+      let(:squash_path) { File.join(repo_path, "gitlab-worktree", "squash-#{subject.id}") }
 
-    before do
-      system(*%W(#{Gitlab.config.git.bin_path} -C #{repo_path} worktree add --detach #{squash_path} master))
-    end
+      before do
+        system(*%W(#{Gitlab.config.git.bin_path} -C #{repo_path} worktree add --detach #{squash_path} master))
+      end
 
-    it 'returns true when there is a current squash directory' do
-      expect(subject.squash_in_progress?).to be_truthy
-    end
+      it 'returns true when there is a current squash directory' do
+        expect(subject.squash_in_progress?).to be_truthy
+      end
 
-    it 'returns false when there is no squash directory' do
-      FileUtils.rm_rf(squash_path)
+      it 'returns false when there is no squash directory' do
+        FileUtils.rm_rf(squash_path)
 
-      expect(subject.squash_in_progress?).to be_falsey
-    end
+        expect(subject.squash_in_progress?).to be_falsey
+      end
 
-    it 'returns false when the squash directory has expired' do
-      time = 20.minutes.ago.to_time
-      File.utime(time, time, squash_path)
+      it 'returns false when the squash directory has expired' do
+        time = 20.minutes.ago.to_time
+        File.utime(time, time, squash_path)
+
+        expect(subject.squash_in_progress?).to be_falsey
+      end
+
+      it 'returns false when the source project has been removed' do
+        allow(subject).to receive(:source_project).and_return(nil)
 
-      expect(subject.squash_in_progress?).to be_falsey
+        expect(subject.squash_in_progress?).to be_falsey
+      end
     end
 
-    it 'returns false when the source project has been removed' do
-      allow(subject).to receive(:source_project).and_return(nil)
+    context 'when Gitaly squash_in_progress is enabled' do
+      it_behaves_like 'checking whether a squash is in progress'
+    end
 
-      expect(subject.squash_in_progress?).to be_falsey
+    context 'when Gitaly squash_in_progress is disabled', :disable_gitaly do
+      it_behaves_like 'checking whether a squash is in progress'
     end
   end
 
@@ -2577,6 +2587,14 @@ describe MergeRequest do
         expect(subject.rebase_in_progress?).to be_falsey
       end
     end
+
+    context 'when Gitaly rebase_in_progress is enabled' do
+      it_behaves_like 'checking whether a rebase is in progress'
+    end
+
+    context 'when Gitaly rebase_in_progress is enabled', :disable_gitaly do
+      it_behaves_like 'checking whether a rebase is in progress'
+    end
   end
 
   describe '#allow_collaboration' do

spec/support/gitaly.rb

+RSpec.configure do |config|
+  config.before(:each) do |example|
+    if example.metadata[:disable_gitaly]
+      # Use 'and_wrap_original' to make sure the arguments are valid
+      allow(Gitlab::GitalyClient).to receive(:feature_enabled?).and_wrap_original { |m, *args| m.call(*args) && false }
+    else
+      next if example.metadata[:skip_gitaly_mock]
+
+      # Use 'and_wrap_original' to make sure the arguments are valid
+      allow(Gitlab::GitalyClient).to receive(:feature_enabled?).and_wrap_original do |m, *args|
+        m.call(*args)
+        !Gitlab::GitalyClient.explicit_opt_in_required.include?(args.first)
+      end
+    end
+  end
+end