Separate GRPC channels per repository storage

4 files changed, 28 insertions, 25 deletions

diff --git a/lib/api/internal.rb b/lib/api/internal.rb
index 7eed93aba00..523f38d129e 100644
--- a/lib/api/internal.rb
+++ b/lib/api/internal.rb
@@ -139,7 +139,7 @@ module API
         return unless Gitlab::GitalyClient.enabled?
 
         begin
-          Gitlab::GitalyClient::Notifications.new.post_receive(params[:repo_path])
+          Gitlab::GitalyClient::Notifications.new(params[:repo_path]).post_receive
         rescue GRPC::Unavailable => e
           render_api_error(e, 500)
         end
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)
diff --git a/lib/gitlab/gitaly_client/commit.rb b/lib/gitlab/gitaly_client/commit.rb
index 525b8d680e9..9c714a3ee45 100644
--- a/lib/gitlab/gitaly_client/commit.rb
+++ b/lib/gitlab/gitaly_client/commit.rb
@@ -7,8 +7,10 @@ module Gitlab
 
       class << self
         def diff_from_parent(commit, options = {})
-          stub      = Gitaly::Diff::Stub.new(nil, nil, channel_override: GitalyClient.channel)
-          repo      = Gitaly::Repository.new(path: commit.project.repository.path_to_repo)
+          project   = commit.project
+          channel   = GitalyClient.get_channel(project.repository_storage)
+          stub      = Gitaly::Diff::Stub.new(nil, nil, channel_override: channel)
+          repo      = Gitaly::Repository.new(path: project.repository.path_to_repo)
           parent    = commit.parents[0]
           parent_id = parent ? parent.id : EMPTY_TREE_ID
           request   = Gitaly::CommitDiffRequest.new(
diff --git a/lib/gitlab/gitaly_client/notifications.rb b/lib/gitlab/gitaly_client/notifications.rb
index b827a56207f..cbfb129c002 100644
--- a/lib/gitlab/gitaly_client/notifications.rb
+++ b/lib/gitlab/gitaly_client/notifications.rb
@@ -3,14 +3,19 @@ module Gitlab
     class Notifications
       attr_accessor :stub
 
-      def initialize
-        @stub = Gitaly::Notifications::Stub.new(nil, nil, channel_override: GitalyClient.channel)
+      def initialize(repo_path)
+        full_path = Gitlab::RepoPath.strip_storage_path(repo_path).
+          sub(/\.git\z/, '').sub(/\.wiki\z/, '')
+        @project = Project.find_by_full_path(full_path)
+
+        channel = GitalyClient.get_channel(@project.repository_storage)
+        @stub = Gitaly::Notifications::Stub.new(nil, nil, channel_override: channel)
       end
 
-      def post_receive(repo_path)
-        repository = Gitaly::Repository.new(path: repo_path)
+      def post_receive
+        repository = Gitaly::Repository.new(path: @project.repository.path_to_repo)
         request = Gitaly::PostReceiveRequest.new(repository: repository)
-        stub.post_receive(request)
+        @stub.post_receive(request)
       end
     end
   end