Multi-Level Step Functions Trace Merging Instructions (#26424)

Co-authored-by: cswatt <cecilia.watt@datadoghq.com>

content/en/serverless/step_functions/merge-step-functions-lambda.md

@@ -8,6 +8,50 @@ further_reading:
 
 This page describes how to merge your AWS Step Functions traces with related AWS Lambda traces or nested Step Functions traces. These instructions assume that you have already instrumented these [AWS Step Functions][1] and [Lambda functions][2] to send traces to Datadog.
 
+<div class="alert alert-info">Datadog recommends using <code>JSONata</code> to define your Step Function payloads for the most complete end-to-end tracing experience. If you are using <code>JSONPath</code> to define your Step Function payloads, see the below sections for supported configurations</a>.</div>
+
+## Merge upstream traces with Step Functions and downstream Lambda traces
+
+### Requirements
+Node.js (layer v116+) or Python (layer v103+) runtimes.
+
+Your State Machine Definition must be using [JSONata][1] as the query language. This can be enabled by setting `"QueryLanguage": "JSONata"` at the top-level of the State Machine Definition.
+
+### Setup
+
+On the Lambda Task, set the `Arguments` key as follows: 
+
+```json
+"Arguments": {
+  "Payload": "{% ($execInput := $states.context.Execution.Input; $hasDatadogTraceId := $exists($execInput._datadog.`x-datadog-trace-id`); $hasDatadogRootExecutionId := $exists($execInput._datadog.RootExecutionId); $ddTraceContext := $hasDatadogTraceId ? {'x-datadog-trace-id': $execInput._datadog.`x-datadog-trace-id`, 'x-datadog-tags': $execInput._datadog.`x-datadog-tags`} : {'RootExecutionId': $hasDatadogRootExecutionId ? $execInput._datadog.RootExecutionId : $states.context.Execution.Id}; $sfnContext := $merge([$states.context, {'Execution': $sift($states.context.Execution, function($v, $k) { $k != 'Input' })}]); $merge([$states.input, {'_datadog': $merge([$sfnContext, $ddTraceContext, {'serverless-version': 'v1'}])}])) %}",
+  ...
+}
+```
+
+The `JSONata` expression merges the upstream service's context with the current Step Functions context object and the Lambda state's input payload.
+
+Alternatively, if you have business logic defined in the payload, you can replace `$states.input` at the end of the `JSONata` expression with your intended value for the `Payload` key.
+
+[1]: https://aws.amazon.com/blogs/compute/simplifying-developer-experience-with-variables-and-jsonata-in-aws-step-functions/
+
+## Merge upstream traces with Step Functions and nested Step Functions traces
+
+### Requirements
+Your State Machine Definition must be using `JSONata` as the query language. This can be enabled by setting `"QueryLanguage": "JSONata"` at the top-level of the State Machine Definition.
+
+### Setup
+
+On the Step Functions Task, set the `_datadog` field in the `Input` key as follows: 
+
+```json
+"Arguments": {
+  "Input": {
+    "_datadog": "{% ($execInput := $states.context.Execution.Input; $hasDatadogTraceId := $exists($execInput._datadog.`x-datadog-trace-id`); $hasDatadogRootExecutionId := $exists($execInput._datadog.RootExecutionId); $ddTraceContext := $hasDatadogTraceId ? {'x-datadog-trace-id': $execInput._datadog.`x-datadog-trace-id`, 'x-datadog-tags': $execInput._datadog.`x-datadog-tags`} : {'RootExecutionId': $hasDatadogRootExecutionId ?  $execInput._datadog.RootExecutionId : $states.context.Execution.Id}; $sfnContext := $merge([$states.context, {'Execution': $sift($states.context.Execution, function($v, $k) { $k != 'Input' })}]); $merge([$sfnContext, $ddTraceContext, {'serverless-version': 'v1'}])) %}",
+    ...
+  }
+}
+```
+
 ## Merge Step Functions traces with downstream Lambda traces
 
 ### Requirements
@@ -178,9 +222,5 @@ To link your Step Function traces to nested Step Function traces, configure your
 }
 {{< /highlight >}}
 
-## Merge Lambda -> Step Functions -> Lambda
-
-This capability is not supported.
-
 [1]: /serverless/step_functions/installation
 [2]: /serverless/installation/#installation-instructions