Inheriting from listener started trace with distributed tracing

[drvn-mr]

Following the documentation and using Datadog::Lambda.wrap, we always seem to have a separate aws.lambda trace.

When the Lambda function is called with distributed tracing headers, the trace continues as expected.
However the first Datadog::Tracing.trace call does not become a child trace of the aws.lambda trace, instead being a direct child of the original continued trace ID.
So in Datadog we end up seeing 2 separate traces.

Is it possible to inherit from the trace started by the listener on_start ?
It doesn't seem to be the active trace.

Expected Behavior

Being able to see a single parent trace when a Lambda function is invoked.

Actual Behavior

With distributed tracing, traces started after Datadog::Lambda.wrap are always separate from the wrapping trace.

Steps to Reproduce the Problem

1. Invoke a Lambda function with event object containing "headers" and injected Datadog headers.
   In our case we just do something like this:

Datadog::Tracing.trace(trace_name) do |_span, trace|
  headers = {}
  Datadog::Tracing::Contrib::HTTP.inject(trace.to_digest, headers)

  client = Aws::Lambda::Client.new
  client.invoke({
    ...
    payload: JSON.generate({ headers: headers, ... })
  })
end

2. In the Lambda handler, wrap and start a trace:

Datadog::Lambda.wrap(event, context) do
  Datadog::Tracing.trace("do the thing") do
    ...
  end

  { statusCode: 200, body: 'Hello World' }
end

Specifications

• Datadog Lambda Layer version: 3.27.0
• Datadog Lambda extension v92
• Ruby version: 3.4.8

[purple4reina]

I suspect what is going on is that you've got the tracer installed twice. Since the tracer is being provided by the layer (assuming you're using the layer) you won't need to include it in your bundled dependencies.

Can you confirm how you're bundling the tracer and if you're using the Datadog Ruby Lambda Layer?

[drvn-mr]

I'm not sure what you mean by "bundled dependencies"..
As I said I've followed the instructions for a Container image Lambda function here: https://docs.datadoghq.com/serverless/aws_lambda/instrumentation/ruby/?tab=containerimage

And it says to install both the gems and the Lambda extension.
So both gems are installed in Gemfile

gem "datadog"
gem "datadog-lambda"

In the Dockerfile, we added COPY --from=public.ecr.aws/datadog/lambda-extension:latest /opt/. /opt/

And the Lambda function handler does nothing more than wrapping everything in Datadog::Lambda.wrap(event, context).
So I assume this is what you mean by "using the Datadog Ruby Lambda Layer" (it is this repo after all).

From what I understand of the code, the issue is that the listener does not set its own trace as "active".
And there's no way to continue from it as it's inaccessible.

[purple4reina]

Ah, my apologies, thank you for the correction. And to confirm, how are you initializing your instrumentation? Here's an example ruby function that I've successfully see distributed tracing work. Wondering what is different between yours and mine. One difference is that mine isn't containerized. Anything else. especially in differences between how you're adding instrumentation?

# client.rb

require 'aws-sdk-sqs'
require 'aws-sdk-sns'
require 'aws-sdk-eventbridge'
require 'aws-sdk-lambda'
require 'datadog/lambda'
require 'net/http'
require 'json'
require 'openssl'    

Datadog::Lambda.configure_apm do |c|
# Enable the instrumentation
end

# must be required after all other requires
require 'datadog/auto_instrument'

$sqs_client = Aws::SQS::Client.new
$sns_client = Aws::SNS::Client.new
$event_bridge_client = Aws::EventBridge::Client.new
$lambda_client = Aws::Lambda::Client.new

def handler(event:, context:)
  Datadog::Lambda.wrap(event, context) do
    config = JSON.parse(event['body'])

    config['url'].each do |url|
      puts "propagating to url transport #{url}"
      Net::HTTP.get(URI(url))
    end

    config['sqs'].each do |url|
      puts "propagating to sqs transport #{url}"
      $sqs_client.send_message(queue_url: url, message_body: 'hello')
    end

    config['sns'].each do |url|
      puts "propagating to sns transport #{url}"
      $sns_client.publish(topic_arn: url, message: 'hello')
    end

    config['eventbridge'].each do |bus_name|
      puts "propagating to eventbridge transport #{bus_name}"
      $event_bridge_client.putevents({
        entries: [{
          event_bus_name: bus_name,
          source: 'trace-propagation.client',
          detail_type: 'trace-propagation-test',
          detail: '{"foo": "bar"}'
        }]
      })
    end

    config['invoke'].each do |invoke_func|
      puts "propagating to lambda invoke transport #{invoke_func}"
      $lambda_client.invoke({
        function_name: invoke_func,
        invocation_type: 'RequestResponse',
        payload: '{"hello": "world"}',
      })
    end

    config['appsync'].each do |api|
      attrs = api.split("   ", -1)
      puts "propagating to appsync transport #{attrs.first}"
      uri = URI.parse(attrs[1])
      http = Net::HTTP.new(uri.host, uri.port)
      http.use_ssl = true
      request = Net::HTTP::Post.new(uri.request_uri)
      params = {'query': 'query { getItems }', 'variables': {} }
      request.body = params.to_json
      request['Content-Type'] = 'application/json'
      request['x-api-key'] = attrs[2]
      res = http.request(request)
    end

    'ok'
  end
end

# server.rb

require 'datadog/lambda'

Datadog::Lambda.configure_apm do |c|
# Enable the instrumentation
end

def handler(event:, context:)
  Datadog::Lambda.wrap(event, context) do
    sleep 0.1
    'ok'
  end
end

One thing we do see with ruby trace propagation often is that sometimes the downstream spans will get orphaned. Which is different than receiving a different trace id. Can you confirm that you are indeed seeing two completely separate traces?

Looking forward to getting to the bottom of this.

[drvn-mr]

As a note, I see your example doesn't inject "headers" into the function payload, so are the headers actually propagated to the Lambda function?
I think it's necessary to do when invoking the function directly without an API gateway unless the client is already some AWS special service.
Which is why I specified that we manually inject the propagation headers in the payload, otherwise dd-trace-rb doesn't do it automatically.

The client (also Ruby) config looks like this:

require "datadog/auto_instrument"

Datadog.configure do |c|
  c.tracing.enabled = true
  c.tracing.partial_flush.enabled = true
  c.tracing.partial_flush.min_spans_threshold = 500

  c.tracing.log_injection = true

  c.tracing.instrument :rails
  c.tracing.instrument :sidekiq, distributed_tracing: true
  c.tracing.instrument :pg
  c.tracing.instrument :http
  c.tracing.instrument :faraday
  c.tracing.instrument :aws
  c.tracing.instrument :action_mailer
end

This is the kind of trace we get with "distributed tracing" enabled (adding "headers") to the payload:

We can see the pdf-merge-dev (function name) trace separate when I'd expect the perform trace to be beneath it.

And when "headers" is not part of the payload, we do have expected output that all traces are under the aws.lambda one.
However the upstream service loses correlation.

[drvn-mr]

I managed to have a working solution with this change drvn-mr@836dce1
And in the Lambda handler:

Datadog::Lambda.wrap(event, context) do
  span = Datadog::Lambda.listener_span
  digest = Datadog::Tracing.active_trace&.to_digest
  if span.present? && digest.present?
    new_digest = digest.merge(
      span_id: span.id,
      span_name: span.name,
      span_resource: span.resource,
      span_service: span.service,
      span_type: span.type,
      trace_id: span.trace_id
    )

    Datadog::Tracing.continue_trace!(new_digest)
  end

  Datadog::Tracing.trace("perform") do
    # do the thing
  end
end

Hopefully this explains what I'm trying to do.
Maybe there's already a way to do this but I haven't figured it out.