Add a gRPC health check to ensure Gitaly is upsh-add-gitaly-health-check

This will enable Geo to skip shards that not operational.

Relates to gitlab-org/gitlab-ee#4329

3 files changed, 90 insertions, 3 deletions

diff --git a/lib/gitlab/gitaly_client.rb b/lib/gitlab/gitaly_client.rb
index 4507ea923b4..6bd256f57c7 100644
--- a/lib/gitlab/gitaly_client.rb
+++ b/lib/gitlab/gitaly_client.rb
@@ -1,6 +1,8 @@
 require 'base64'
 
 require 'gitaly'
+require 'grpc/health/v1/health_pb'
+require 'grpc/health/v1/health_services_pb'
 
 module Gitlab
   module GitalyClient
@@ -69,14 +71,27 @@ module Gitlab
         @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 = stub_class(name)
+          addr = stub_address(storage)
           klass.new(addr, :this_channel_is_insecure)
         end
       end
     end
 
+    def self.stub_class(name)
+      if name == :health_check
+        Grpc::Health::V1::Health::Stub
+      else
+        Gitaly.const_get(name.to_s.camelcase.to_sym).const_get(:Stub)
+      end
+    end
+
+    def self.stub_address(storage)
+      addr = address(storage)
+      addr = addr.sub(%r{^tcp://}, '') if URI(addr).scheme == 'tcp'
+      addr
+    end
+
     def self.clear_stubs!
       MUTEX.synchronize do
         @stubs = nil
diff --git a/lib/gitlab/gitaly_client/health_check_service.rb b/lib/gitlab/gitaly_client/health_check_service.rb
new file mode 100644
index 00000000000..6c1213f5e20
--- /dev/null
+++ b/lib/gitlab/gitaly_client/health_check_service.rb
@@ -0,0 +1,19 @@
+module Gitlab
+  module GitalyClient
+    class HealthCheckService
+      def initialize(storage)
+        @storage = storage
+      end
+
+      # Sends a gRPC health ping to the Gitaly server for the storage shard.
+      def check
+        request = Grpc::Health::V1::HealthCheckRequest.new
+        response = GitalyClient.call(@storage, :health_check, :check, request, timeout: GitalyClient.fast_timeout)
+
+        { success: response&.status == :SERVING }
+      rescue GRPC::BadStatus => e
+        { success: false, message: e.to_s }
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/health_checks/gitaly_check.rb b/lib/gitlab/health_checks/gitaly_check.rb
new file mode 100644
index 00000000000..11416c002e3
--- /dev/null
+++ b/lib/gitlab/health_checks/gitaly_check.rb
@@ -0,0 +1,53 @@
+module Gitlab
+  module HealthChecks
+    class GitalyCheck
+      extend BaseAbstractCheck
+
+      METRIC_PREFIX = 'gitaly_health_check'.freeze
+
+      class << self
+        def readiness
+          repository_storages.map do |storage_name|
+            check(storage_name)
+          end
+        end
+
+        def metrics
+          repository_storages.flat_map do |storage_name|
+            result, elapsed = with_timing { check(storage_name) }
+            labels = { shard: storage_name }
+
+            [
+              metric("#{metric_prefix}_success", successful?(result) ? 1 : 0, **labels),
+              metric("#{metric_prefix}_latency_seconds", elapsed, **labels)
+            ].flatten
+          end
+        end
+
+        def check(storage_name)
+          serv = Gitlab::GitalyClient::HealthCheckService.new(storage_name)
+          result = serv.check
+          HealthChecks::Result.new(result[:success], result[:message], shard: storage_name)
+        end
+
+        private
+
+        def metric_prefix
+          METRIC_PREFIX
+        end
+
+        def successful?(result)
+          result[:success]
+        end
+
+        def repository_storages
+          storages.keys
+        end
+
+        def storages
+          Gitlab.config.repositories.storages
+        end
+      end
+    end
+  end
+end