Factor out 'temporarily allow' logicgitaly-tripswitch

2 files changed, 41 insertions, 25 deletions

diff --git a/lib/gitlab/gitaly_client/storage_settings.rb b/lib/gitlab/gitaly_client/storage_settings.rb
index 8e8d6c18519..02fcb413abd 100644
--- a/lib/gitlab/gitaly_client/storage_settings.rb
+++ b/lib/gitlab/gitaly_client/storage_settings.rb
@@ -4,6 +4,8 @@ module Gitlab
     # where production code (app, config, db, lib) touches Git repositories
     # directly.
     class StorageSettings
+      extend Gitlab::TemporarilyAllow
+
       DirectPathAccessError = Class.new(StandardError)
       InvalidConfigurationError = Class.new(StandardError)
 
@@ -20,37 +22,15 @@ module Gitlab
       MUTEX = Mutex.new
 
       DISK_ACCESS_DENIED_FLAG = :deny_disk_access
+      ALLOW_KEY = :allow_disk_access
 
       # If your code needs this method then your code needs to be fixed.
       def self.allow_disk_access
-        self.disk_access_override_add(1)
-        yield
-      ensure
-        self.disk_access_override_add(-1)
-      end
-
-      def self.disk_access_override?
-        if RequestStore.active?
-          RequestStore.fetch(:gitaly_disk_access_override, 0) > 0
-        else
-          MUTEX.synchronize { @disk_access_override && @disk_access_override > 0 }
-        end
-      end
-
-      def self.disk_access_override_add(value)
-        if RequestStore.active?
-          RequestStore[:gitaly_disk_access_override] ||= 0
-          RequestStore[:gitaly_disk_access_override] += value
-        else
-          MUTEX.synchronize do
-            @disk_access_override ||= 0
-            @disk_access_override += value
-          end
-        end
+        temporarily_allow(ALLOW_KEY) { yield }
       end
 
       def self.disk_access_denied?
-        !disk_access_override? && GitalyClient.feature_enabled?(DISK_ACCESS_DENIED_FLAG)
+        !temporarily_allowed?(ALLOW_KEY) && GitalyClient.feature_enabled?(DISK_ACCESS_DENIED_FLAG)
       rescue
         false # Err on the side of caution, don't break gitlab for people
       end
diff --git a/lib/gitlab/temporarily_allow.rb b/lib/gitlab/temporarily_allow.rb
new file mode 100644
index 00000000000..2fd2fc63617
--- /dev/null
+++ b/lib/gitlab/temporarily_allow.rb
@@ -0,0 +1,36 @@
+module Gitlab
+  module TemporarilyAllow
+    TEMPORARILY_ALLOW_MUTEX = Mutex.new
+
+    def temporarily_allow(key)
+      temporarily_allow_add(key, 1)
+      yield
+    ensure
+      temporarily_allow_add(key, -1)
+    end
+
+    def temporarily_allowed?(key)
+      if RequestStore.active?
+        h = RequestStore[:temporarily_allow]
+        h && h[key] > 0
+      else
+        TEMPORARILY_ALLOW_MUTEX.synchronize do
+          h = @temporarily_allow
+          h && h[key] > 0
+        end
+      end
+    end
+
+    def temporarily_allow_add(key, value)
+      if RequestStore.active?
+        RequestStore[:temporarily_allow] ||= Hash.new(0)
+        RequestStore[:temporarily_allow][key] += value
+      else
+        TEMPORARILY_ALLOW_MUTEX.synchronize do
+          @temporarily_allow ||= Hash.new(0)
+          @temporarily_allow[key] += value
+        end
+      end
+    end
+  end
+end