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, 123 insertions, 0 deletions

diff --git a/spec/lib/gitlab/gitaly_client/health_check_service_spec.rb b/spec/lib/gitlab/gitaly_client/health_check_service_spec.rb
new file mode 100644
index 00000000000..2c7e5eb5787
--- /dev/null
+++ b/spec/lib/gitlab/gitaly_client/health_check_service_spec.rb
@@ -0,0 +1,41 @@
+require 'spec_helper'
+
+describe Gitlab::GitalyClient::HealthCheckService do
+  let(:project) { create(:project) }
+  let(:storage_name) { project.repository_storage }
+
+  subject { described_class.new(storage_name) }
+
+  describe '#check' do
+    it 'successfully sends a health check request' do
+      expect(Gitlab::GitalyClient).to receive(:call).with(
+        storage_name,
+        :health_check,
+        :check,
+        instance_of(Grpc::Health::V1::HealthCheckRequest),
+        timeout: Gitlab::GitalyClient.fast_timeout).and_call_original
+
+      expect(subject.check).to eq({ success: true })
+    end
+
+    it 'receives an unsuccessful health check request' do
+      expect_any_instance_of(Grpc::Health::V1::Health::Stub)
+        .to receive(:check)
+        .and_return(double(status: false))
+
+      expect(subject.check).to eq({ success: false })
+    end
+
+    it 'gracefully handles gRPC error' do
+      expect(Gitlab::GitalyClient).to receive(:call).with(
+        storage_name,
+        :health_check,
+        :check,
+        instance_of(Grpc::Health::V1::HealthCheckRequest),
+        timeout: Gitlab::GitalyClient.fast_timeout)
+          .and_raise(GRPC::Unavailable.new('Connection refused'))
+
+      expect(subject.check).to eq({ success: false, message: '14:Connection refused' })
+    end
+  end
+end
diff --git a/spec/lib/gitlab/gitaly_client_spec.rb b/spec/lib/gitlab/gitaly_client_spec.rb
index 309b7338ef0..81bcd8c28ed 100644
--- a/spec/lib/gitlab/gitaly_client_spec.rb
+++ b/spec/lib/gitlab/gitaly_client_spec.rb
@@ -3,6 +3,31 @@ require 'spec_helper'
 # We stub Gitaly in `spec/support/gitaly.rb` for other tests. We don't want
 # those stubs while testing the GitalyClient itself.
 describe Gitlab::GitalyClient, skip_gitaly_mock: true do
+  describe '.stub_class' do
+    it 'returns the gRPC health check stub' do
+      expect(described_class.stub_class(:health_check)).to eq(::Grpc::Health::V1::Health::Stub)
+    end
+
+    it 'returns a Gitaly stub' do
+      expect(described_class.stub_class(:ref_service)).to eq(::Gitaly::RefService::Stub)
+    end
+  end
+
+  describe '.stub_address' do
+    it 'returns the same result after being called multiple times' do
+      address = 'localhost:9876'
+      prefixed_address = "tcp://#{address}"
+
+      allow(Gitlab.config.repositories).to receive(:storages).and_return({
+        'default' => { 'gitaly_address' => prefixed_address }
+      })
+
+      2.times do
+        expect(described_class.stub_address('default')).to eq('localhost:9876')
+      end
+    end
+  end
+
   describe '.stub' do
     # Notice that this is referring to gRPC "stubs", not rspec stubs
     before do
diff --git a/spec/lib/gitlab/health_checks/gitaly_check_spec.rb b/spec/lib/gitlab/health_checks/gitaly_check_spec.rb
new file mode 100644
index 00000000000..724beefff69
--- /dev/null
+++ b/spec/lib/gitlab/health_checks/gitaly_check_spec.rb
@@ -0,0 +1,57 @@
+require 'spec_helper'
+
+describe Gitlab::HealthChecks::GitalyCheck do
+  let(:result_class) { Gitlab::HealthChecks::Result }
+  let(:repository_storages) { ['default'] }
+
+  before do
+    allow(described_class).to receive(:repository_storages) { repository_storages }
+  end
+
+  describe '#readiness' do
+    subject { described_class.readiness }
+
+    before do
+      expect(Gitlab::GitalyClient::HealthCheckService).to receive(:new).and_return(gitaly_check)
+    end
+
+    context 'Gitaly server is up' do
+      let(:gitaly_check) { double(check: { success: true }) }
+
+      it { is_expected.to eq([result_class.new(true, nil, shard: 'default')]) }
+    end
+
+    context 'Gitaly server is down' do
+      let(:gitaly_check) { double(check: { success: false, message: 'Connection refused' }) }
+
+      it { is_expected.to eq([result_class.new(false, 'Connection refused', shard: 'default')]) }
+    end
+  end
+
+  describe '#metrics' do
+    subject { described_class.metrics }
+
+    before do
+      expect(Gitlab::GitalyClient::HealthCheckService).to receive(:new).and_return(gitaly_check)
+    end
+
+    context 'Gitaly server is up' do
+      let(:gitaly_check) { double(check: { success: true }) }
+
+      it 'provides metrics' do
+        expect(subject).to all(have_attributes(labels: { shard: 'default' }))
+        expect(subject).to include(an_object_having_attributes(name: 'gitaly_health_check_success', value: 1))
+        expect(subject).to include(an_object_having_attributes(name: 'gitaly_health_check_latency_seconds', value: be >= 0))
+      end
+    end
+
+    context 'Gitaly server is down' do
+      let(:gitaly_check) { double(check: { success: false, message: 'Connection refused' }) }
+
+      it 'provides metrics' do
+        expect(subject).to include(an_object_having_attributes(name: 'gitaly_health_check_success', value: 0))
+        expect(subject).to include(an_object_having_attributes(name: 'gitaly_health_check_latency_seconds', value: be >= 0))
+      end
+    end
+  end
+end