1 files changed, 9 insertions, 3 deletions
diff --git a/lib/gitlab/database/load_balancing/sticking.rb b/lib/gitlab/database/load_balancing/sticking.rb
index 8e5dc98e96e..f5cb83e398a 100644
--- a/lib/gitlab/database/load_balancing/sticking.rb
+++ b/lib/gitlab/database/load_balancing/sticking.rb
@@ -121,19 +121,19 @@ module Gitlab
         end
 
         def unstick(namespace, id)
-          Gitlab::Redis::SharedState.with do |redis|
+          with_redis do |redis|
             redis.del(redis_key_for(namespace, id))
           end
         end
 
         def set_write_location_for(namespace, id, location)
-          Gitlab::Redis::SharedState.with do |redis|
+          with_redis do |redis|
             redis.set(redis_key_for(namespace, id), location, ex: EXPIRATION)
           end
         end
 
         def last_write_location_for(namespace, id)
-          Gitlab::Redis::SharedState.with do |redis|
+          with_redis do |redis|
             redis.get(redis_key_for(namespace, id))
           end
         end
@@ -143,6 +143,12 @@ module Gitlab
 
           "database-load-balancing/write-location/#{name}/#{namespace}/#{id}"
         end
+
+        private
+
+        def with_redis(&block)
+          Gitlab::Redis::DbLoadBalancing.with(&block)
+        end
       end
     end
   end