Adds a GRPC client interceptor for CorrelationIDs (8c568f01) · Commits · GitLab.org / ruby / gems / labkit-ruby

gitlab-labkit.gemspec

+1 −1

Original line number	Diff line number	Diff line
		@@ -30,7 +30,7 @@ Gem::Specification.new do \|spec\|
		spec.add_development_dependency "grpc-tools", "~> 1.19"
		spec.add_development_dependency "rack", "~> 2.0"
		spec.add_development_dependency "rake", "~> 12.3"
		spec.add_development_dependency "rspec", "~> 3.6.0"
		spec.add_development_dependency "rspec", "~> 3.8.0"
		spec.add_development_dependency "rspec-parameterized", "~> 0.4"
		spec.add_development_dependency "rubocop", "~> 0.65.0"
		spec.add_development_dependency "rubocop-rspec", "~> 1.22.1"

lib/labkit/correlation/grpc.rb

+1 −0

Original line number	Diff line number	Diff line
		@@ -4,6 +4,7 @@ module Labkit
		module Correlation
		# The GRPC module contains functionality for instrumenting GRPC calls
		module GRPC
		autoload :ClientInterceptor, "labkit/correlation/grpc/client_interceptor"
		autoload :GRPCCommon, "labkit/correlation/grpc/grpc_common"
		autoload :ServerInterceptor, "labkit/correlation/grpc/server_interceptor"
		end

lib/labkit/correlation/grpc/client_interceptor.rb

0 → 100644

+52 −0

Original line number	Diff line number	Diff line
		# frozen_string_literal: true

		# Disable the UnusedMethodArgument linter, since we need to declare the kwargs
		# in the methods, but we don't actually use them.
		# rubocop:disable Lint/UnusedMethodArgument

		require "grpc"

		module Labkit
		module Correlation
		module GRPC
		# ClientInterceptor is used to inject the correlation_id into the metadata
		# or a GRPC call for onward propagation to the server
		class ClientInterceptor < ::GRPC::ClientInterceptor
		include Labkit::Correlation::GRPC::GRPCCommon
		include Singleton

		def request_response(request:, call:, method:, metadata:)
		inject_correlation_id_into_metadata(metadata)

		yield
		end

		def client_streamer(requests:, call:, method:, metadata:)
		inject_correlation_id_into_metadata(metadata)

		yield
		end

		def server_streamer(request:, call:, method:, metadata:)
		inject_correlation_id_into_metadata(metadata)

		yield
		end

		def bidi_streamer(requests:, call:, method:, metadata:)
		inject_correlation_id_into_metadata(metadata)

		yield
		end

		private

		def inject_correlation_id_into_metadata(metadata, &block)
		metadata[CORRELATION_METADATA_KEY] = Labkit::Correlation::CorrelationId.current_id if Labkit::Correlation::CorrelationId.current_id
		end
		end
		end
		end
		end

		# rubocop:enable Lint/UnusedMethodArgument

spec/labkit/correlation/grpc/client_interceptor_spec.rb

0 → 100644

+134 −0

Original line number	Diff line number	Diff line
		# frozen_string_literal: true

		require "./spec/support/grpc_service"

		describe Labkit::Correlation::GRPC::ClientInterceptor do
		describe "Labkit::Correlation::GRPC::ClientInterceptor" do
		subject { described_class.instance }

		shared_examples_for "a grpc interceptor method" do
		let(:custom_error) { Class.new(StandardError) }

		it "yields" do
		expect { \|b\| method.call(kwargs, &b) }.to yield_control
		end

		it "propagates exceptions" do
		expect { method.call(kwargs) { raise custom_error } }.to raise_error(custom_error)
		end
		end

		describe "#request_response" do
		let(:method) { subject.method(:request_response) }
		let(:kwargs) { { request: {}, call: {}, method: "grc_method", metadata: {} } }

		it_behaves_like "a grpc interceptor method"
		end

		describe "#client_streamer" do
		let(:method) { subject.method(:client_streamer) }
		let(:kwargs) { { requests: [], call: {}, method: "grc_method", metadata: {} } }

		it_behaves_like "a grpc interceptor method"
		end

		describe "#server_streamer" do
		let(:method) { subject.method(:server_streamer) }
		let(:kwargs) { { request: {}, call: {}, method: "grc_method", metadata: {} } }

		it_behaves_like "a grpc interceptor method"
		end

		describe "#bidi_streamer" do
		let(:method) { subject.method(:bidi_streamer) }
		let(:kwargs) { { requests: [], call: {}, method: "grc_method", metadata: {} } }

		it_behaves_like "a grpc interceptor method"
		end
		end

		describe "running on RpcServer" do
		let(:mock_server) { LabkitTest::TestService::MockServer.new }
		let(:client) { mock_server.start(client_interceptors: [described_class.instance]) }
		let(:correlation_id) { "123456" }

		after do
		mock_server.stop
		end

		RSpec::Matchers.define :metadata_includes_correlation_id do \|correlation_id\|
		match { \|actual\| actual.metadata[Labkit::Correlation::GRPC::GRPCCommon::CORRELATION_METADATA_KEY] == correlation_id }
		end

		RSpec::Matchers.define :metadata_does_not_include_correlation_id do
		match { \|actual\| !actual.metadata.include?(Labkit::Correlation::GRPC::GRPCCommon::CORRELATION_METADATA_KEY) }
		end

		describe "#req_res_method" do
		it "sends correlation_id when available" do
		expect(Labkit::Correlation::CorrelationId).to receive(:current_id).at_least(:once).and_return(correlation_id)
		expect(mock_server.server_impl).to receive(:req_res_method).with(anything, metadata_includes_correlation_id(correlation_id)).and_call_original

		expect(client.req_res_method(LabkitTest::Msg.new)).to be_a(LabkitTest::Msg)
		end

		it "does not send correlation_id when unavailable" do
		expect(Labkit::Correlation::CorrelationId).to receive(:current_id).at_least(:once).and_return(nil)
		expect(mock_server.server_impl).to receive(:req_res_method).with(anything, metadata_does_not_include_correlation_id).and_call_original

		expect(client.req_res_method(LabkitTest::Msg.new)).to be_a(LabkitTest::Msg)
		end
		end

		describe "#server_stream_method" do
		it "sends correlation_id when available" do
		expect(Labkit::Correlation::CorrelationId).to receive(:current_id).at_least(:once).and_return(correlation_id)
		expect(mock_server.server_impl).to receive(:server_stream_method).with(anything, metadata_includes_correlation_id(correlation_id)).and_call_original

		expect(client.server_stream_method(LabkitTest::Msg.new)).to all(be_a(LabkitTest::Msg))
		end

		it "does not send correlation_id when unavailable" do
		expect(Labkit::Correlation::CorrelationId).to receive(:current_id).at_least(:once).and_return(nil)
		expect(mock_server.server_impl).to receive(:server_stream_method).with(anything, metadata_does_not_include_correlation_id).and_call_original

		expect(client.server_stream_method(LabkitTest::Msg.new)).to all(be_a(LabkitTest::Msg))
		end
		end

		describe "#client_stream_method" do
		it "sends correlation_id when available" do
		expect(Labkit::Correlation::CorrelationId).to receive(:current_id).at_least(:once).and_return(correlation_id)
		expect(mock_server.server_impl).to receive(:client_stream_method).with(metadata_includes_correlation_id(correlation_id)).and_call_original

		client.client_stream_method([LabkitTest::Msg.new])
		end

		it "does not send correlation_id when unavailable" do
		expect(Labkit::Correlation::CorrelationId).to receive(:current_id).at_least(:once).and_return(nil)
		expect(mock_server.server_impl).to receive(:client_stream_method).with(metadata_does_not_include_correlation_id).and_call_original

		client.client_stream_method([LabkitTest::Msg.new])
		end
		end

		# Note that in the `bidi_stream_method` we stub `bidi_stream_method_shadow`, not `bidi_stream_method`.
		# This is due to a bad interaction between rspec and GRPC wrt to function arity.
		# Stubbing the shadow method is a workaround.
		describe "#bidi_stream_method" do
		it "sends correlation_id when available" do
		expect(Labkit::Correlation::CorrelationId).to receive(:current_id).at_least(:once).and_return(correlation_id)
		expect(mock_server.server_impl).to receive(:bidi_stream_method_shadow).with(anything, metadata_includes_correlation_id(correlation_id)).and_call_original

		expect(client.bidi_stream_method([LabkitTest::Msg.new])).to all(be_a(LabkitTest::Msg))
		end

		it "does not send correlation_id when unavailable" do
		expect(Labkit::Correlation::CorrelationId).to receive(:current_id).at_least(:once).and_return(nil)
		expect(mock_server.server_impl).to receive(:bidi_stream_method_shadow).with(anything, metadata_does_not_include_correlation_id).and_call_original

		expect(client.bidi_stream_method([LabkitTest::Msg.new])).to all(be_a(LabkitTest::Msg))
		end
		end
		end
		end

spec/labkit/correlation/grpc/server_interceptor_spec.rb

+1 −1

Original line number	Diff line number	Diff line
		@@ -51,7 +51,7 @@ describe Labkit::Correlation::GRPC::ServerInterceptor do
		expect(Labkit::Correlation::CorrelationId).to receive(:use_id).with(expected_correlation_id).and_call_original

		enumerator = client.bidi_stream_method([LabkitTest::Msg.new], metadata: metadata_with_correlation_id)
		enumerator.each { } # Consume the stream
		enumerator.each { \|x\| } # Consume the stream
		end
		end
		end