1 files changed, 31 insertions, 15 deletions
diff --git a/lib/gitlab/gitaly_client.rb b/lib/gitlab/gitaly_client.rb
index a0dbe0a8c11..72466700c05 100644
--- a/lib/gitlab/gitaly_client.rb
+++ b/lib/gitlab/gitaly_client.rb
@@ -4,26 +4,42 @@ module Gitlab
   module GitalyClient
     SERVER_VERSION_FILE = 'GITALY_SERVER_VERSION'.freeze
 
-    def self.configure_channel(storage, address)
-      @addresses ||= {}
-      @addresses[storage] = address
-      @channels ||= {}
-      @channels[storage] = new_channel(address)
-    end
+    MUTEX = Mutex.new
+    private_constant :MUTEX
 
-    def self.new_channel(address)
-      # NOTE: Gitaly currently runs on a Unix socket, so permissions are
-      # handled using the file system and no additional authentication is
-      # required (therefore the :this_channel_is_insecure flag)
-      GRPC::Core::Channel.new(address, {}, :this_channel_is_insecure)
+    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
+      end
     end
 
-    def self.get_channel(storage)
-      @channels[storage]
+    def self.clear_stubs!
+      MUTEX.synchronize do
+        @stubs = nil
+      end
     end
 
-    def self.get_address(storage)
-      @addresses[storage]
+    def self.address(storage)
+      params = Gitlab.config.repositories.storages[storage]
+      raise "storage not found: #{storage.inspect}" if params.nil?
+
+      address = params['gitaly_address']
+      unless address.present?
+        raise "storage #{storage.inspect} is missing a gitaly_address"
+      end
+
+      unless URI(address).scheme.in?(%w(tcp unix))
+        raise "Unsupported Gitaly address: #{address.inspect} does not use URL scheme 'tcp' or 'unix'"
+      end
+
+      address
     end
 
     def self.enabled?