Separate GRPC channels per repository storage

1 files changed, 13 insertions, 17 deletions

diff --git a/lib/gitlab/gitaly_client.rb b/lib/gitlab/gitaly_client.rb
index 1ce47ef2b05..c947075bf62 100644
--- a/lib/gitlab/gitaly_client.rb
+++ b/lib/gitlab/gitaly_client.rb
@@ -4,28 +4,24 @@ module Gitlab
   module GitalyClient
     SERVER_VERSION_FILE = 'GITALY_SERVER_VERSION'.freeze
 
-    def self.gitaly_address
-      if Gitlab.config.gitaly.socket_path
-        "unix://#{Gitlab.config.gitaly.socket_path}"
-      end
+    def self.configure_channel(shard, socket_path)
+      @channel ||= {}
+      @channel[shard] = new_channel("unix://#{socket_path}")
+    end
+
+    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)
     end
 
-    def self.channel
-      return @channel if defined?(@channel)
-
-      @channel =
-        if enabled?
-          # 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(gitaly_address, {}, :this_channel_is_insecure)
-        else
-          nil
-        end
+    def self.get_channel(shard)
+      @channel.fetch(shard)
     end
 
     def self.enabled?
-      gitaly_address.present?
+      Gitlab.config.gitaly.enabled
     end
 
     def self.feature_enabled?(feature)