Use `Feature` gates on `Gitlab::GitalyClient.feature_enabled?`feature-flags-flipper

4 files changed, 110 insertions, 4 deletions

diff --git a/lib/feature.rb b/lib/feature.rb
index b89e5176a08..2e2b343f82c 100644
--- a/lib/feature.rb
+++ b/lib/feature.rb
@@ -20,6 +20,13 @@ class Feature
       flipper.feature(key)
     end
 
+    def persisted?(feature)
+      # Flipper creates on-memory features when asked for a not-yet-created one.
+      # If we want to check if a feature has been actually set, we look for it
+      # on the persisted features list.
+      all.map(&:name).include?(feature.name)
+    end
+
     private
 
     def flipper
diff --git a/lib/gitlab/gitaly_client.rb b/lib/gitlab/gitaly_client.rb
index 72466700c05..2343446bf22 100644
--- a/lib/gitlab/gitaly_client.rb
+++ b/lib/gitlab/gitaly_client.rb
@@ -2,6 +2,12 @@ require 'gitaly'
 
 module Gitlab
   module GitalyClient
+    module MigrationStatus
+      DISABLED = 1
+      OPT_IN = 2
+      OPT_OUT = 3
+    end
+
     SERVER_VERSION_FILE = 'GITALY_SERVER_VERSION'.freeze
 
     MUTEX = Mutex.new
@@ -46,8 +52,20 @@ module Gitlab
       Gitlab.config.gitaly.enabled
     end
 
-    def self.feature_enabled?(feature)
-      enabled? && ENV["GITALY_#{feature.upcase}"] == '1'
+    def self.feature_enabled?(feature, status: MigrationStatus::OPT_IN)
+      return false if !enabled? || status == MigrationStatus::DISABLED
+
+      feature = Feature.get("gitaly_#{feature}")
+
+      # If the feature hasn't been set, turn it on if it's opt-out
+      return status == MigrationStatus::OPT_OUT unless Feature.persisted?(feature)
+
+      if feature.percentage_of_time_value > 0
+        # Probabilistically enable this feature
+        return Random.rand() * 100 < feature.percentage_of_time_value
+      end
+
+      feature.enabled?
     end
 
     def self.migrate(feature)
diff --git a/spec/lib/gitlab/gitaly_client_spec.rb b/spec/lib/gitlab/gitaly_client_spec.rb
index 08ee0dff6b2..95ecba67532 100644
--- a/spec/lib/gitlab/gitaly_client_spec.rb
+++ b/spec/lib/gitlab/gitaly_client_spec.rb
@@ -1,7 +1,10 @@
 require 'spec_helper'
 
-describe Gitlab::GitalyClient, lib: true do
+# 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, lib: true, skip_gitaly_mock: true do
   describe '.stub' do
+    # Notice that this is referring to gRPC "stubs", not rspec stubs
     before { described_class.clear_stubs! }
 
     context 'when passed a UNIX socket address' do
@@ -32,4 +35,81 @@ describe Gitlab::GitalyClient, lib: true do
       end
     end
   end
+
+  describe 'feature_enabled?' do
+    let(:feature_name) { 'my_feature' }
+    let(:real_feature_name) { "gitaly_#{feature_name}" }
+
+    context 'when Gitaly is disabled' do
+      before { allow(described_class).to receive(:enabled?).and_return(false) }
+
+      it 'returns false' do
+        expect(described_class.feature_enabled?(feature_name)).to be(false)
+      end
+    end
+
+    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 { Feature.get(real_feature_name).disable }
+
+        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 { Feature.get(real_feature_name).enable }
+
+        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 { Feature.get(real_feature_name).enable_percentage_of_time(70) }
+
+        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
+    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 { Feature.get(real_feature_name).disable }
+
+        it 'returns false' do
+          expect(described_class.feature_enabled?(feature_name, status: feature_status)).to be(false)
+        end
+      end
+    end
+  end
 end
diff --git a/spec/support/gitaly.rb b/spec/support/gitaly.rb
index 7aca902fc61..2bf159002a0 100644
--- a/spec/support/gitaly.rb
+++ b/spec/support/gitaly.rb
@@ -1,6 +1,7 @@
 if Gitlab::GitalyClient.enabled?
   RSpec.configure do |config|
-    config.before(:each) do
+    config.before(:each) do |example|
+      next if example.metadata[:skip_gitaly_mock]
       allow(Gitlab::GitalyClient).to receive(:feature_enabled?).and_return(true)
     end
   end