From 97e931d1941108844aacd5a66b1e6f4e602bc4a2 Mon Sep 17 00:00:00 2001
From: Nao Minami <south37777@gmail.com>
Date: Mon, 20 Jan 2020 18:38:30 +0900
Subject: [PATCH 1/2] Add gRPC interceptor

---
 .../integrations/grpc_server_interceptor.rb   | 213 +++++++++++++++++
 opencensus.gemspec                            |   1 +
 .../grpc_server_interceptor_test.rb           | 219 ++++++++++++++++++
 3 files changed, 433 insertions(+)
 create mode 100644 lib/opencensus/trace/integrations/grpc_server_interceptor.rb
 create mode 100644 test/trace/integrations/grpc_server_interceptor_test.rb

diff --git a/lib/opencensus/trace/integrations/grpc_server_interceptor.rb b/lib/opencensus/trace/integrations/grpc_server_interceptor.rb
new file mode 100644
index 00000000..ed8fabd7
--- /dev/null
+++ b/lib/opencensus/trace/integrations/grpc_server_interceptor.rb
@@ -0,0 +1,213 @@
+require "opencensus"
+
+module OpenCensus
+  module Trace
+    module Integrations
+      ##
+      # # gRPC interceptor
+      #
+      # This is a interceptor for gRPC:
+      #
+      # * It wraps all incoming requests in a root span
+      # * It exports the captured spans at the end of the request.
+      #
+      # Example:
+      #
+      #     require "opencensus/trace/integrations/grpc_server_interceptor"
+      #
+      #     server = GRPC::RpcServer.new(
+      #       interceptors: [
+      #         OpenCensus::Trace::Integrations::GrpcServerInterceptor.new,
+      #       ]
+      #     )
+      #
+      class GrpcServerInterceptor
+        ##
+        # A key we use to read the parent span context.
+        #
+        # @private
+        #
+        OPENCENSUS_TRACE_BIN_KEY = "grpc-trace-bin".freeze
+
+        ##
+        # @param [#export] exporter The exported used to export captured spans
+        #     at the end of the request. Optional: If omitted, uses the exporter
+        #     in the current config.
+        # @param [#call] span_modifier Modify span if necessary. It takes span,
+        # request, call, method as its parameters.
+        #
+        def initialize exporter: nil, span_modifier: nil
+          @exporter = exporter || OpenCensus::Trace.config.exporter
+          @span_modifier = span_modifier
+          @formatter = Formatters::Binary.new
+        end
+
+        ##
+        # Intercept a unary request response call.
+        #
+        # @param [Object] request
+        # @param [GRPC::ActiveCall::SingleReqView] call
+        # @param [Method] method
+        #
+        def request_response request:, call:, method:
+          context_bin = call.metadata[OPENCENSUS_TRACE_BIN_KEY]
+          if context_bin
+            context = deserialize(context_bin)
+          end
+
+          Trace.start_request_trace \
+            trace_context: context,
+            same_process_as_parent: false do |span_context|
+            begin
+              Trace.in_span get_name(method) do |span|
+                modify_span span, request, call, method
+
+                start_request span, call, method
+                begin
+                  grpc_ex = GRPC::Ok.new
+                  yield
+                rescue StandardError => e
+                  grpc_ex = to_grpc_ex(e)
+                  raise e
+                ensure
+                  finish_request span, grpc_ex
+                end
+              end
+            ensure
+              @exporter.export span_context.build_contained_spans
+            end
+          end
+        end
+
+        # NOTE: For now, we don't support server_streamer, client_streamer and
+        # bidi_streamer
+
+        private
+
+        ##
+        # @param [String] context_bin OpenCensus span context in binary format
+        # @return [OpenCensus::Trace::TraceContextData, nil]
+        #
+        def deserialize context_bin
+          @formatter.deserialize(context_bin)
+        end
+
+        ##
+        # Span name is represented as $package.$service/$method
+        # cf. https://github.com/census-instrumentation/opencensus-specs/blob/master/trace/gRPC.md#spans
+        #
+        # @param [Method] method
+        # @return [String]
+        #
+        def get_name method
+          "#{method.owner.service_name}/#{camelize(method.name.to_s)}"
+        end
+
+        ##
+        # @param [Method] method
+        # @return [String]
+        #
+        def get_path method
+          "/" + get_name(method)
+        end
+
+        ##
+        # @param [String] term
+        # @return [String]
+        #
+        def camelize term
+          term.split("_").map(&:capitalize).join
+        end
+
+        ##
+        # Modify span by custom span modifier
+        #
+        # @param [OpenCensus::Trace::SpanBuilder] span
+        # @param [Object] request
+        # @param [GRPC::ActiveCall::SingleReqView] call
+        # @param [Method] method
+        #
+        def modify_span span, request, call, method
+          @span_modifier.call(span, request, call, method) if @span_modifier
+        end
+
+        ##
+        # @param [OpenCensus::Trace::SpanBuilder] span
+        # @param [GRPC::ActiveCall::SingleReqView] call
+        # @param [Method] method
+        #
+        def start_request span, call, method
+          span.kind = SpanBuilder::SERVER
+          span.put_attribute "http.path", get_path(method)
+          span.put_attribute "http.method", "POST"  # gRPC always uses "POST"
+          if call.metadata['user-agent']
+            span.put_attribute "http.user_agent", call.metadata['user-agent']
+          end
+        end
+
+        ##
+        # @param [OpenCensus::Trace::SpanBuilder] span
+        # @param [GRPC::BadStatus] exception
+        #
+        def finish_request span, exception
+          # Set gRPC server status
+          # https://github.com/census-instrumentation/opencensus-specs/blob/master/trace/gRPC.md#spans
+          span.set_status exception.code
+          span.put_attribute "http.status_code", to_http_status(exception)
+        end
+
+        ##
+        # cf. https://github.com/census-instrumentation/opencensus-specs/blob/master/trace/HTTP.md#mapping-from-http-status-codes-to-trace-status-codes
+        #
+        # @param [GRPC::BadStatus] exception
+        # @return [Integer]
+        #
+        def to_http_status exception
+          case exception
+          when GRPC::Ok
+            200
+          when GRPC::InvalidArgument
+            400
+          when GRPC::DeadlineExceeded
+            504
+          when GRPC::NotFound
+            404
+          when GRPC::PermissionDenied
+            403
+          when GRPC::Unauthenticated
+            401
+          when GRPC::Aborted
+            # For GRPC::Aborted, grpc-gateway uses 409. We do the same.
+            # cf. https://github.com/grpc-ecosystem/grpc-gateway/blob/e8db07a3923d3f5c77dbcea96656afe43a2757a8/runtime/errors.go#L17-L58
+            409
+          when GRPC::ResourceExhausted
+            429
+          when GRPC::Unimplemented
+            501
+          when GRPC::Unavailable
+            503
+          when GRPC::Unknown
+            # NOTE: This is not same with the correct mapping
+            500
+          else
+            # NOTE: Here, we use 500 temporarily.
+            500
+          end
+        end
+
+        ##
+        # @param [Exception] exception
+        # @return [GRPC::BadStatus]
+        #
+        def to_grpc_ex exception
+          case exception
+          when GRPC::BadStatus
+            exception
+          else
+            GRPC::Unknown.new(exception.message)
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/opencensus.gemspec b/opencensus.gemspec
index 890cdb1a..e16bfbc8 100644
--- a/opencensus.gemspec
+++ b/opencensus.gemspec
@@ -29,4 +29,5 @@ Gem::Specification.new do |spec|
   spec.add_development_dependency "rubocop", "~> 0.59.2"
   spec.add_development_dependency "yard", "~> 0.9"
   spec.add_development_dependency "yard-doctest", "~> 0.1.6"
+  spec.add_development_dependency "grpc", "~> 1.26"
 end
diff --git a/test/trace/integrations/grpc_server_interceptor_test.rb b/test/trace/integrations/grpc_server_interceptor_test.rb
new file mode 100644
index 00000000..c39d9aff
--- /dev/null
+++ b/test/trace/integrations/grpc_server_interceptor_test.rb
@@ -0,0 +1,219 @@
+require "test_helper"
+
+require "google/protobuf/empty_pb"
+require "google/protobuf/wrappers_pb"
+require "grpc"
+require "opencensus/trace/integrations/grpc_server_interceptor"
+
+describe OpenCensus::Trace::Integrations::GrpcServerInterceptor do
+  module TestRpc
+    class Service
+      include GRPC::GenericService
+
+      self.marshal_class_method = :encode
+      self.unmarshal_class_method = :decode
+      self.service_name = "test.TestRpc"
+
+      rpc :HelloRpc, Google::Protobuf::StringValue, Google::Protobuf::Empty
+    end
+  end
+
+  class TestService < TestRpc::Service
+    def hello_rpc(req, call)
+      # Do nothing
+    end
+  end
+
+  class MockedExporter
+    attr_reader :spans
+
+    def initialize
+      @spans = []
+    end
+
+    def export spans
+      @spans += spans
+    end
+  end
+
+  MockedCall = Struct.new(:metadata)
+
+  let(:exporter) { MockedExporter.new }
+  let(:request) { Google::Protobuf::StringValue.new(value: "World") }
+  let(:call) { MockedCall.new(metadata) }
+  let(:method_object) { TestService.new.method(:hello_rpc) }
+
+  describe "basic request" do
+    let(:interceptor) { OpenCensus::Trace::Integrations::GrpcServerInterceptor.new exporter: exporter }
+    let(:metadata) {
+      {
+        "user-agent" => "Google Chrome".encode(Encoding::ASCII_8BIT),
+      }
+    }
+    let(:spans) do
+      # make sure the request is processed
+      interceptor.request_response(
+        request: request, call: call, method: method_object) {}
+      exporter.spans
+    end
+    let(:root_span) { spans.first }
+
+    it "captures spans" do
+      spans.wont_be_empty
+      spans.count.must_equal 1
+    end
+
+    it "parses the request path" do
+      root_span.name.value.must_equal "test.TestRpc/HelloRpc"
+    end
+
+    it "captures the response status code" do
+      root_span.status.wont_be_nil
+      root_span.status.code.must_equal OpenCensus::Trace::Status::OK
+    end
+
+    it "adds attributes to the span" do
+      root_span.kind.must_equal :SERVER
+      root_span.attributes["http.method"].value.must_equal "POST"
+      root_span.attributes["http.path"].value.must_equal "/test.TestRpc/HelloRpc"
+      root_span.attributes["http.user_agent"].value.must_equal "Google Chrome"
+    end
+  end
+
+  describe "failure response" do
+    let(:interceptor) { OpenCensus::Trace::Integrations::GrpcServerInterceptor.new exporter: exporter }
+    let(:metadata) { {} }
+
+    it "captures the response status code" do
+      assert_raises GRPC::NotFound do
+        interceptor.request_response \
+          request: request,
+          call: call,
+          method: method_object do
+          raise GRPC::NotFound
+        end
+      end
+      root_span = exporter.spans.first
+
+      root_span.status.wont_be_nil
+      root_span.status.code.must_equal OpenCensus::Trace::Status::NOT_FOUND
+    end
+  end
+
+  describe "global configuration" do
+    let(:metadata) { {} }
+
+    describe "default exporter" do
+      let(:interceptor) { OpenCensus::Trace::Integrations::GrpcServerInterceptor.new  }
+
+      it "should use the default Logger exporter" do
+        out, _err = capture_subprocess_io do
+          interceptor.request_response(
+            request: request, call: call, method: method_object) {}
+        end
+        out.wont_be_empty
+      end
+    end
+
+    describe "custom exporter" do
+      before do
+        @original_exporter = OpenCensus::Trace.config.exporter
+        OpenCensus::Trace.config.exporter = exporter
+      end
+      after do
+        OpenCensus::Trace.config.exporter = @original_exporter
+      end
+      let(:interceptor) { OpenCensus::Trace::Integrations::GrpcServerInterceptor.new  }
+
+      it "should capture the request" do
+        interceptor.request_response(
+          request: request, call: call, method: method_object) {}
+
+        spans = exporter.spans
+        spans.wont_be_empty
+        spans.count.must_equal 1
+      end
+    end
+  end
+
+  describe "trace context formatting" do
+    let(:interceptor) { OpenCensus::Trace::Integrations::GrpcServerInterceptor.new exporter: exporter }
+
+    describe "metadata with valid grpc-trace-bin" do
+      let(:metadata) {
+        {
+          "grpc-trace-bin" => OpenCensus::Trace::Formatters::Binary.new.serialize(span_context),
+        }
+      }
+      let(:span_context) {
+        OpenCensus::Trace::TraceContextData.new(
+          "0123456789abcdef0123456789abcdef",  # trace_id
+          "0123456789abcdef",                  # span_id
+          0,                                   # trace_options
+        )
+      }
+
+      it "parses trace-context header from rack environment" do
+        interceptor.request_response(
+          request: request, call: call, method: method_object) {}
+        root_span = exporter.spans.first
+
+        root_span.trace_id.must_equal "0123456789abcdef0123456789abcdef"
+        root_span.parent_span_id.must_equal "0123456789abcdef"
+      end
+    end
+
+    describe "metadata with missing grpc-trace-bin" do
+      let(:metadata) { {} }
+
+      it "falls back to default for missing header" do
+        interceptor.request_response(
+          request: request, call: call, method: method_object) {}
+        root_span = exporter.spans.first
+
+        root_span.trace_id.must_match %r{^[0-9a-f]{32}$}
+        root_span.parent_span_id.must_be_empty
+      end
+    end
+
+    describe "metadata with invalid grpc-trace-bin" do
+      let(:metadata) {
+        {
+          "grpc-trace-bin" => "invalid".encode(Encoding::ASCII_8BIT),
+        }
+      }
+
+      it "falls back to default for invalid trace-context version" do
+        interceptor.request_response(
+          request: request, call: call, method: method_object) {}
+        root_span = exporter.spans.first
+
+        root_span.trace_id.must_match %r{^[0-9a-f]{32}$}
+        root_span.parent_span_id.must_be_empty
+      end
+    end
+  end
+
+  describe "span modifier" do
+    let(:interceptor) {
+      OpenCensus::Trace::Integrations::GrpcServerInterceptor.new(
+        exporter: exporter,
+        span_modifier: span_modifier
+      )
+    }
+    let(:span_modifier) {
+      -> (span, request, call, method) {
+        span.put_attribute "test.test_attribute", "dummy-value"
+      }
+    }
+    let(:metadata) { {} }
+
+    it "modifies a root span" do
+      interceptor.request_response(
+        request: request, call: call, method: method_object) {}
+      root_span = exporter.spans.first
+
+      root_span.attributes["test.test_attribute"].value.must_equal "dummy-value"
+    end
+  end
+end

From fd45541deb3852bf33d5bac170ec26f71d7194c5 Mon Sep 17 00:00:00 2001
From: Nao Minami <south37777@gmail.com>
Date: Mon, 20 Jan 2020 21:24:44 +0900
Subject: [PATCH 2/2] Fix rubocop violations

---
 .../integrations/grpc_server_interceptor.rb   | 55 ++++++++++++-------
 1 file changed, 34 insertions(+), 21 deletions(-)

diff --git a/lib/opencensus/trace/integrations/grpc_server_interceptor.rb b/lib/opencensus/trace/integrations/grpc_server_interceptor.rb
index ed8fabd7..5b07f001 100644
--- a/lib/opencensus/trace/integrations/grpc_server_interceptor.rb
+++ b/lib/opencensus/trace/integrations/grpc_server_interceptor.rb
@@ -49,30 +49,15 @@ def initialize exporter: nil, span_modifier: nil
         # @param [GRPC::ActiveCall::SingleReqView] call
         # @param [Method] method
         #
-        def request_response request:, call:, method:
+        def request_response request:, call:, method:, &block
           context_bin = call.metadata[OPENCENSUS_TRACE_BIN_KEY]
-          if context_bin
-            context = deserialize(context_bin)
-          end
+          context = context_bin ? deserialize(context_bin) : nil
 
           Trace.start_request_trace \
             trace_context: context,
             same_process_as_parent: false do |span_context|
             begin
-              Trace.in_span get_name(method) do |span|
-                modify_span span, request, call, method
-
-                start_request span, call, method
-                begin
-                  grpc_ex = GRPC::Ok.new
-                  yield
-                rescue StandardError => e
-                  grpc_ex = to_grpc_ex(e)
-                  raise e
-                ensure
-                  finish_request span, grpc_ex
-                end
-              end
+              yield_with_trace(request, call, method, &block)
             ensure
               @exporter.export span_context.build_contained_spans
             end
@@ -92,6 +77,28 @@ def deserialize context_bin
           @formatter.deserialize(context_bin)
         end
 
+        ##
+        # @param [Object] request
+        # @param [GRPC::ActiveCall::SingleReqView] call
+        # @param [Method] method
+        #
+        def yield_with_trace request, call, method
+          Trace.in_span get_name(method) do |span|
+            modify_span span, request, call, method
+
+            start_request span, call, method
+            begin
+              grpc_ex = GRPC::Ok.new
+              yield request: request, call: call, method: method
+            rescue StandardError => e
+              grpc_ex = to_grpc_ex(e)
+              raise e
+            ensure
+              finish_request span, grpc_ex
+            end
+          end
+        end
+
         ##
         # Span name is represented as $package.$service/$method
         # cf. https://github.com/census-instrumentation/opencensus-specs/blob/master/trace/gRPC.md#spans
@@ -139,9 +146,9 @@ def modify_span span, request, call, method
         def start_request span, call, method
           span.kind = SpanBuilder::SERVER
           span.put_attribute "http.path", get_path(method)
-          span.put_attribute "http.method", "POST"  # gRPC always uses "POST"
-          if call.metadata['user-agent']
-            span.put_attribute "http.user_agent", call.metadata['user-agent']
+          span.put_attribute "http.method", "POST" # gRPC always uses "POST"
+          if call.metadata["user-agent"]
+            span.put_attribute "http.user_agent", call.metadata["user-agent"]
           end
         end
 
@@ -156,6 +163,9 @@ def finish_request span, exception
           span.put_attribute "http.status_code", to_http_status(exception)
         end
 
+        # rubocop:disable Metrics/MethodLength
+        # rubocop:disable Metrics/CyclomaticComplexity
+
         ##
         # cf. https://github.com/census-instrumentation/opencensus-specs/blob/master/trace/HTTP.md#mapping-from-http-status-codes-to-trace-status-codes
         #
@@ -195,6 +205,9 @@ def to_http_status exception
           end
         end
 
+        # rubocop:enable Metrics/MethodLength
+        # rubocop:enable Metrics/CyclomaticComplexity
+
         ##
         # @param [Exception] exception
         # @return [GRPC::BadStatus]