Add a correlation GRPC server middleware (b1761676) · Commits · GitLab.org / ruby / gems / labkit-ruby

lib/labkit/correlation.rb

+2 −0

Original line number	Diff line number	Diff line
		@@ -3,6 +3,8 @@
		module Labkit
		# Correlation provides correlation functionality
		module Correlation
		autoload :GRPC, "labkit/correlation/grpc"

		autoload :CorrelationId, "labkit/correlation/correlation_id"
		end
		end

lib/labkit/correlation/grpc.rb

0 → 100644

+11 −0

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

		module Labkit
		module Correlation
		# The GRPC module contains functionality for instrumenting GRPC calls
		module GRPC
		autoload :GRPCCommon, "labkit/correlation/grpc/grpc_common"
		autoload :ServerInterceptor, "labkit/correlation/grpc/server_interceptor"
		end
		end
		end

lib/labkit/correlation/grpc/grpc_common.rb

0 → 100644

+13 −0

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

		module Labkit
		module Correlation
		module GRPC
		# This module is shared between the client and server interceptor middlewares.
		# It is not part of the public API
		module GRPCCommon
		CORRELATION_METADATA_KEY = "x-gitlab-correlation-id"
		end
		end
		end
		end

lib/labkit/correlation/grpc/server_interceptor.rb

0 → 100644

+55 −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
		# ServerInterceptor is a server-side GRPC interceptor
		# for injecting GRPC calls with a correlation-id passed from
		# a GRPC client to the GRPC Ruby Service
		class ServerInterceptor < ::GRPC::ServerInterceptor
		include Labkit::Correlation::GRPC::GRPCCommon

		def request_response(request: nil, call: nil, method: nil)
		wrap_with_correlation_id(call) do
		yield
		end
		end

		def client_streamer(call: nil, method: nil)
		wrap_with_correlation_id(call) do
		yield
		end
		end

		def server_streamer(request: nil, call: nil, method: nil)
		wrap_with_correlation_id(call) do
		yield
		end
		end

		def bidi_streamer(requests: nil, call: nil, method: nil)
		wrap_with_correlation_id(call) do
		yield
		end
		end

		private

		def wrap_with_correlation_id(call, &block)
		correlation_id = call.metadata[CORRELATION_METADATA_KEY]
		correlation_id \|\|= Labkit::Correlation::CorrelationId.current_or_new_id

		Labkit::Correlation::CorrelationId.use_id(correlation_id, &block)
		end
		end
		end
		end
		end

		# rubocop:enable Lint/UnusedMethodArgument

spec/labkit/correlation/grpc/server_interceptor_spec.rb

0 → 100644

+73 −0

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

		require "./spec/support/grpc_service"

		describe Labkit::Correlation::GRPC::ServerInterceptor do
		describe "running on RpcServer" do
		subject { described_class.new }

		let(:mock_server) { LabkitTest::TestService::MockServer.new }
		let(:client) { mock_server.start(server_interceptors: [subject]) }

		after do
		mock_server.stop
		end

		shared_examples_for "a server interceptor" do
		def metadata_with_correlation_id
		tags = {}
		tags[Labkit::Correlation::GRPC::GRPCCommon::CORRELATION_METADATA_KEY] = correlation_id if correlation_id

		tags
		end

		describe "#req_res_method" do
		it "generates instrumentation" do
		expect(Labkit::Correlation::CorrelationId).to receive(:use_id).with(expected_correlation_id).and_call_original

		client.req_res_method(LabkitTest::Msg.new, metadata: metadata_with_correlation_id)
		end
		end

		describe "#server_stream_method" do
		it "generates instrumentation" do
		expect(Labkit::Correlation::CorrelationId).to receive(:use_id).with(expected_correlation_id).and_call_original

		enumerator = client.server_stream_method(LabkitTest::Msg.new, metadata: metadata_with_correlation_id)
		enumerator.each { } # Consume the stream
		end
		end

		describe "#client_stream_method" do
		it "generates instrumentation" do
		expect(Labkit::Correlation::CorrelationId).to receive(:use_id).with(expected_correlation_id).and_call_original

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

		describe "#bidi_stream_method" do
		it "generates instrumentation" 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
		end
		end
		end

		describe "with correlation_id" do
		let(:correlation_id) { "12345" }
		let(:expected_correlation_id) { correlation_id }

		it_behaves_like "a server interceptor"
		end

		describe "without correlation_id" do
		let(:correlation_id) { nil }
		let(:expected_correlation_id) { anything }

		it_behaves_like "a server interceptor"
		end
		end
		end