1 files changed, 25 insertions, 39 deletions

diff --git a/lib/gitlab/gitaly_client.rb b/lib/gitlab/gitaly_client.rb
index c69676a1dac..72466700c05 100644
--- a/lib/gitlab/gitaly_client.rb
+++ b/lib/gitlab/gitaly_client.rb
@@ -4,56 +4,42 @@ module Gitlab
   module GitalyClient
     SERVER_VERSION_FILE = 'GITALY_SERVER_VERSION'.freeze
 
-    # This function is not thread-safe because it updates Hashes in instance variables.
-    def self.configure_channels
-      @addresses = {}
-      @channels = {}
-      Gitlab.config.repositories.storages.each do |name, params|
-        address = params['gitaly_address']
-        unless address.present?
-          raise "storage #{name.inspect} is missing a gitaly_address"
-        end
+    MUTEX = Mutex.new
+    private_constant :MUTEX
 
-        unless URI(address).scheme.in?(%w(tcp unix))
-          raise "Unsupported Gitaly address: #{address.inspect} does not use URL scheme 'tcp' or 'unix'"
+    def self.stub(name, storage)
+      MUTEX.synchronize do
+        @stubs ||= {}
+        @stubs[storage] ||= {}
+        @stubs[storage][name] ||= begin
+          klass = Gitaly.const_get(name.to_s.camelcase.to_sym).const_get(:Stub)
+          addr = address(storage)
+          addr = addr.sub(%r{^tcp://}, '') if URI(addr).scheme == 'tcp'
+          klass.new(addr, :this_channel_is_insecure)
         end
-
-        @addresses[name] = address
-        @channels[name] = new_channel(address)
       end
     end
 
-    def self.new_channel(address)
-      address = address.sub(%r{^tcp://}, '') if URI(address).scheme == 'tcp'
-      # NOTE: When Gitaly runs on a Unix socket, permissions are
-      # handled using the file system and no additional authentication is
-      # required (therefore the :this_channel_is_insecure flag)
-      # TODO: Add authentication support when Gitaly is running on a TCP socket.
-      GRPC::Core::Channel.new(address, {}, :this_channel_is_insecure)
+    def self.clear_stubs!
+      MUTEX.synchronize do
+        @stubs = nil
+      end
     end
 
-    def self.get_channel(storage)
-      if !Rails.env.production? && @channels.nil?
-        # In development mode the Rails auto-loader may reset the instance
-        # variables of this class. What we do here is not thread-safe. In normal
-        # circumstances (including production) these instance variables have
-        # been initialized from config/initializers.
-        configure_channels
-      end
+    def self.address(storage)
+      params = Gitlab.config.repositories.storages[storage]
+      raise "storage not found: #{storage.inspect}" if params.nil?
 
-      @channels[storage]
-    end
+      address = params['gitaly_address']
+      unless address.present?
+        raise "storage #{storage.inspect} is missing a gitaly_address"
+      end
 
-    def self.get_address(storage)
-      if !Rails.env.production? && @addresses.nil?
-        # In development mode the Rails auto-loader may reset the instance
-        # variables of this class. What we do here is not thread-safe. In normal
-        # circumstances (including development) these instance variables have
-        # been initialized from config/initializers.
-        configure_channels
+      unless URI(address).scheme.in?(%w(tcp unix))
+        raise "Unsupported Gitaly address: #{address.inspect} does not use URL scheme 'tcp' or 'unix'"
       end
 
-      @addresses[storage]
+      address
     end
 
     def self.enabled?