Commit 73333c29 authored by Jacob Vosmaer's avatar Jacob Vosmaer 👋
Browse files

Create gRPC logging middleware

parent 04c8fb01

lib/labkit/correlation/grpc/grpc_common.rb

+10 −0
Original line number Diff line number Diff line
@@ -7,6 +7,16 @@ module Labkit 
      # It is not part of the public API

      module GRPCCommon

        CORRELATION_METADATA_KEY = "x-gitlab-correlation-id"



        def rpc_split(method)

          owner = method.owner

          method_name, = owner.rpc_descs.find do |k, _|

            ::GRPC::GenericService.underscore(k.to_s) == method.name.to_s

          end

          method_name ||= "(unknown)"



          [owner.service_name, method_name]

        end

      end

    end

  end

lib/labkit/logging.rb

+1 −0
Original line number Diff line number Diff line
@@ -4,6 +4,7 @@ module Labkit 
  # Logging provides functionality for logging, such as

  # sanitization

  module Logging

    autoload :GRPC, "labkit/logging/grpc"

    autoload :Sanitizer, "labkit/logging/sanitizer"

  end

end

lib/labkit/logging/grpc.rb

0 → 100644
+9 −0
Original line number Diff line number Diff line 
# frozen_string_literal: true



module Labkit

  module Logging

    module GRPC

      autoload :ServerInterceptor, "labkit/logging/grpc/server_interceptor"

    end

  end

end

lib/labkit/logging/grpc/server_interceptor.rb

0 → 100644
+84 −0
Original line number Diff line number Diff line 
# frozen_string_literal: true



require "grpc"

require "json"



module Labkit

  module Logging

    module GRPC

      class ServerInterceptor < ::GRPC::ServerInterceptor

        include Labkit::Correlation::GRPC::GRPCCommon



        CODE_STRINGS = {

          ::GRPC::Core::StatusCodes::OK => "OK",

          ::GRPC::Core::StatusCodes::CANCELLED => "Canceled",

          ::GRPC::Core::StatusCodes::UNKNOWN => "Unknown",

          ::GRPC::Core::StatusCodes::INVALID_ARGUMENT => "InvalidArgument",

          ::GRPC::Core::StatusCodes::DEADLINE_EXCEEDED => "DeadlineExceeded",

          ::GRPC::Core::StatusCodes::NOT_FOUND => "NotFound",

          ::GRPC::Core::StatusCodes::ALREADY_EXISTS => "AlreadyExists",

          ::GRPC::Core::StatusCodes::PERMISSION_DENIED => "PermissionDenied",

          ::GRPC::Core::StatusCodes::RESOURCE_EXHAUSTED => "ResourceExhausted",

          ::GRPC::Core::StatusCodes::FAILED_PRECONDITION => "FailedPrecondition",

          ::GRPC::Core::StatusCodes::ABORTED => "Aborted",

          ::GRPC::Core::StatusCodes::OUT_OF_RANGE => "OutOfRange",

          ::GRPC::Core::StatusCodes::UNIMPLEMENTED => "Unimplemented",

          ::GRPC::Core::StatusCodes::INTERNAL => "Internal",

          ::GRPC::Core::StatusCodes::UNAVAILABLE => "Unavailable",

          ::GRPC::Core::StatusCodes::DATA_LOSS => "DataLoss",

          ::GRPC::Core::StatusCodes::UNAUTHENTICATED => "Unauthenticated",

        }.freeze



        def initialize(log_file, default_tags)

          @log_file = log_file

          @log_file.sync = true

          @default_tags = default_tags



          super()

        end



        def request_response(request: nil, call: nil, method: nil)

          log_request(method, call) { yield }

        end



        def server_streamer(request: nil, call: nil, method: nil)

          log_request(method, call) { yield }

        end



        def client_streamer(call: nil, method: nil)

          log_request(method, call) { yield }

        end



        def bidi_streamer(requests: nil, call: nil, method: nil)

          log_request(method, call) { yield }

        end



        private



        def log_request(method, call)

          start = Time.now

          code = ::GRPC::Core::StatusCodes::OK



          yield

        rescue StandardError => ex

          code = ex.is_a?(::GRPC::BadStatus) ? ex.code : ::GRPC::Core::StatusCodes::UNKNOWN



          raise

        ensure

          service_name, method_name = rpc_split(method)

          message = @default_tags.merge(

            'grpc.time_ms': ((Time.now - start) * 1000.0).truncate(3),

            'grpc.code': CODE_STRINGS.fetch(code, code.to_s),

            'grpc.method': method_name,

            'grpc.service': service_name,

            pid: Process.pid,

            correlation_id: call.metadata.fetch(CORRELATION_METADATA_KEY, "(missing)"),

            time: Time.now.utc.strftime("%Y-%m-%dT%H:%M:%S.%LZ"),

          )



          @log_file.puts(JSON.dump(message))

        end

      end

    end

  end

end

lib/labkit/tracing/grpc/server_interceptor.rb

+3 −8
Original line number Diff line number Diff line
@@ -14,6 +14,8 @@ module Labkit 
      # for instrumenting GRPC calls with distributed tracing

      # in a GRPC Ruby server

      class ServerInterceptor < ::GRPC::ServerInterceptor

        include Labkit::Correlation::GRPC::GRPCCommon



        def request_response(request: nil, call: nil, method: nil)

          wrap_with_tracing(call, method, "unary") do

            yield
@@ -40,16 +42,9 @@ module Labkit 


        private



        def route_from_method(method)

          service_class = method.owner

          rpc_method = method.name.to_s.split("_").map(&:capitalize).join("")



          "/#{service_class.service_name}/#{rpc_method}"

        end



        def wrap_with_tracing(call, method, grpc_type)

          context = TracingUtils.tracer.extract(OpenTracing::FORMAT_TEXT_MAP, call.metadata)

          method_name = route_from_method(method)

          method_name = "/#{rpc_split(method).join("/")}"

          tags = {

            "component" => "grpc",

            "span.kind" => "server",