Create proper spans for ActionCable

Author: Taucher2003

app/channels/application_cable/channel.rb

@@ -2,12 +2,20 @@
 
 module ApplicationCable
   class Channel < ActionCable::Channel::Base
+    def self.tracer
+      @tracer ||= ::OpenTelemetry.tracer_provider.tracer('sagittarius-cable')
+    end
+
     def subscribe_to_channel
-      with_context { super }
+      with_otel_span("#{self.class.name} subscribe", 'subscribe') do
+        with_context { super }
+      end
     end
 
     def perform_action(data)
-      with_context { super }
+      with_otel_span("#{self.class.name} #{data['action']}", 'process') do
+        with_context { super }
+      end
     end
 
     protected
@@ -30,6 +38,30 @@ def create_authentication(token_type, token)
       end
     end
 
+    def otel_context
+      return ::OpenTelemetry::Context.current unless connection.respond_to?(:otel_context)
+
+      connection.otel_context || ::OpenTelemetry::Context.current
+    end
+
+    def with_otel_span(span_name, operation, parent_context: otel_context, links: [], &block)
+      attributes = {
+        'messaging.system' => 'action_cable',
+        'messaging.operation' => operation,
+        'code.namespace' => self.class.name,
+      }
+
+      ::OpenTelemetry::Context.with_current(parent_context) do
+        self.class.tracer.in_span(span_name, kind: :server, attributes: attributes, links: links) do |span|
+          block.call(span)
+        rescue StandardError => e
+          span.record_exception(e)
+          span.status = ::OpenTelemetry::Trace::Status.error(e.message)
+          raise
+        end
+      end
+    end
+
     def with_context(&block)
       Code0::ZeroTrack::Context.with_context(
         application: 'cable',

app/channels/application_cable/connection.rb

@@ -2,5 +2,18 @@
 
 module ApplicationCable
   class Connection < ActionCable::Connection::Base
+    attr_reader :otel_context
+
+    def connect
+      @otel_context = extract_otel_context
+    end
+
+    private
+
+    def extract_otel_context
+      ::OpenTelemetry.propagation.extract(request.headers)
+    rescue StandardError
+      ::OpenTelemetry::Context.current
+    end
   end
 end

app/channels/graphql_channel.rb

@@ -11,37 +11,53 @@ def subscribed
   end
 
   def execute(data)
-    result = SagittariusSchema.execute(
-      query: data['query'],
-      context: {
-        current_authentication: find_authentication(@token),
-        visibility_profile: :execution,
-        channel: self,
-      },
-      variables: data['variables'],
-      operation_name: data['operationName']
-    )
-
-    @subscription_ids << result.context[:subscription_id] if result.context[:subscription_id]
-
-    transmit({ result: result.to_h, more: result.subscription? })
+    connection_span_context = ::OpenTelemetry::Trace.current_span(otel_context).context
+    links = if connection_span_context.valid?
+              [::OpenTelemetry::Trace::Link.new(connection_span_context)]
+            else
+              []
+            end
+
+    with_otel_span("#{self.class.name} execute", 'process',
+                   parent_context: ::OpenTelemetry::Context::ROOT, links: links) do |span|
+      span.set_attribute('graphql.operation.name', data['operationName']) if data['operationName']
+
+      result = SagittariusSchema.execute(
+        query: data['query'],
+        context: {
+          current_authentication: find_authentication(@token),
+          visibility_profile: :execution,
+          channel: self,
+        },
+        variables: data['variables'],
+        operation_name: data['operationName']
+      )
+
+      @subscription_ids << result.context[:subscription_id] if result.context[:subscription_id]
+
+      transmit({ result: result.to_h, more: result.subscription? })
+    end
   end
 
   def unsubscribed
-    @subscription_ids.each do |sid|
-      SagittariusSchema.subscriptions.delete_subscription(sid)
+    with_otel_span("#{self.class.name} unsubscribed", 'process') do
+      @subscription_ids.each do |sid|
+        SagittariusSchema.subscriptions.delete_subscription(sid)
+      end
     end
   end
 
   private
 
   def verify_authentication
-    with_context do
-      authentication = find_authentication(@token)
-      return unless authentication.invalid? || authentication.none?
-
-      @subscription_ids.each { |sid| SagittariusSchema.subscriptions.delete_subscription(sid) }
-      reject
+    with_otel_span("#{self.class.name} verify_authentication", 'process') do
+      with_context do
+        authentication = find_authentication(@token)
+        return unless authentication.invalid? || authentication.none?
+
+        @subscription_ids.each { |sid| SagittariusSchema.subscriptions.delete_subscription(sid) }
+        reject
+      end
     end
   end
 end

spec/channels/application_cable/channel_spec.rb

@@ -0,0 +1,102 @@
+# frozen_string_literal: true
+
+require 'rails_helper'
+
+RSpec.describe ApplicationCable::Channel do
+  include AuthenticationHelpers
+  include ActionCable::Channel::TestCase::Behavior
+
+  include_context 'with graphql subscription support'
+
+  tests GraphqlChannel
+
+  let(:exporter) { OpenTelemetry::SDK::Trace::Export::InMemorySpanExporter.new }
+  let(:span_processor) { OpenTelemetry::SDK::Trace::Export::SimpleSpanProcessor.new(exporter) }
+  let(:tracer_provider) do
+    OpenTelemetry::SDK::Trace::TracerProvider.new.tap do |tp|
+      tp.add_span_processor(span_processor)
+    end
+  end
+  let(:tracer) { tracer_provider.tracer('sagittarius-cable') }
+
+  let(:user) { create(:user) }
+  let(:token) { "Session #{authorization_token(user)}" }
+
+  before do
+    propagator = OpenTelemetry::Trace::Propagation::TraceContext::TextMapPropagator.new
+    allow(OpenTelemetry).to receive(:propagation).and_return(propagator)
+    allow(described_class).to receive(:tracer).and_return(tracer)
+    allow(GraphqlChannel).to receive(:tracer).and_return(tracer)
+  end
+
+  describe '#subscribe_to_channel' do
+    it 'creates a span with correct name' do
+      stub_connection(otel_context: OpenTelemetry::Context.current)
+      subscribe(token: token)
+
+      span = exporter.finished_spans.find { |s| s.name.include?('subscribe') }
+      expect(span.name).to eq('GraphqlChannel subscribe')
+    end
+
+    it 'sets messaging attributes' do
+      stub_connection(otel_context: OpenTelemetry::Context.current)
+      subscribe(token: token)
+
+      span = exporter.finished_spans.find { |s| s.name.include?('subscribe') }
+      expect(span.attributes['messaging.system']).to eq('action_cable')
+      expect(span.attributes['messaging.operation']).to eq('subscribe')
+      expect(span.attributes['code.namespace']).to eq('GraphqlChannel')
+    end
+
+    context 'when connection has a trace context' do
+      let(:parent_span_context) do
+        OpenTelemetry::Trace::SpanContext.new(
+          trace_id: OpenTelemetry::Trace.generate_trace_id,
+          span_id: OpenTelemetry::Trace.generate_span_id,
+          trace_flags: OpenTelemetry::Trace::TraceFlags::SAMPLED
+        )
+      end
+      let(:parent_context) do
+        OpenTelemetry::Trace.context_with_span(
+          OpenTelemetry::Trace.non_recording_span(parent_span_context)
+        )
+      end
+
+      it 'creates spans as children of the connection context' do
+        stub_connection(otel_context: parent_context)
+        subscribe(token: token)
+
+        span = exporter.finished_spans.find { |s| s.name.include?('subscribe') }
+        expect(span.hex_trace_id).to eq(parent_span_context.hex_trace_id)
+        expect(span.hex_parent_span_id).to eq(parent_span_context.hex_span_id)
+      end
+    end
+  end
+
+  describe '#perform_action' do
+    before do
+      stub_connection(otel_context: OpenTelemetry::Context.current)
+      subscribe(token: token)
+    end
+
+    it 'creates a span with the action name' do
+      perform :execute,
+              query: 'subscription($message: String) { echo(message: $message) { message } }',
+              variables: { message: 'hello' }
+
+      span = exporter.finished_spans.find { |s| s.name == 'GraphqlChannel execute' }
+      expect(span).not_to be_nil
+    end
+
+    it 'sets messaging attributes' do
+      perform :execute,
+              query: 'subscription($message: String) { echo(message: $message) { message } }',
+              variables: { message: 'hello' }
+
+      span = exporter.finished_spans.find { |s| s.name == 'GraphqlChannel execute' }
+      expect(span).not_to be_nil
+      expect(span.attributes['messaging.system']).to eq('action_cable')
+      expect(span.attributes['code.namespace']).to eq('GraphqlChannel')
+    end
+  end
+end

spec/channels/application_cable/connection_spec.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+require 'rails_helper'
+
+RSpec.describe ApplicationCable::Connection do
+  let(:propagator) { OpenTelemetry::Trace::Propagation::TraceContext::TextMapPropagator.new }
+
+  before do
+    allow(OpenTelemetry).to receive(:propagation).and_return(propagator)
+  end
+
+  describe '#connect' do
+    context 'when traceparent is present in upgrade request headers' do
+      it 'extracts and stores the trace context' do
+        connect '/cable', headers: { 'Traceparent' => '00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01' }
+
+        span_context = OpenTelemetry::Trace.current_span(connection.otel_context).context
+        expect(span_context.hex_trace_id).to eq('0af7651916cd43dd8448eb211c80319c')
+        expect(span_context.hex_span_id).to eq('b7ad6b7169203331')
+      end
+    end
+
+    context 'when no trace headers are present' do
+      it 'stores a context without error' do
+        connect '/cable'
+
+        expect(connection.otel_context).not_to be_nil
+      end
+    end
+
+    context 'when invalid traceparent is present' do
+      it 'stores a context without error' do
+        connect '/cable', headers: { 'Traceparent' => 'garbage-value' }
+
+        expect(connection.otel_context).not_to be_nil
+      end
+    end
+  end
+end