Merge pull request #422 from solarwinds/NH-78538

Add missing SW key propagation

Author: raphael-theriault-swi

.yarn/versions/d01421ba.yml

@@ -0,0 +1,2 @@
+releases:
+  solarwinds-apm: minor

packages/solarwinds-apm/src/propagation/trace-context.ts

@@ -0,0 +1,54 @@
+/*
+Copyright 2023-2024 SolarWinds Worldwide, LLC.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+import {
+  type Context,
+  createTraceState,
+  type SpanContext,
+  type TextMapSetter,
+  trace,
+} from "@opentelemetry/api"
+import { W3CTraceContextPropagator } from "@opentelemetry/core"
+
+const TRACE_STATE_KEY = "tracestate"
+
+export function swValue(context: SpanContext): string {
+  return `${context.spanId}-${context.traceFlags.toString(16).padStart(2, "0")}`
+}
+
+export class TraceContextPropagator extends W3CTraceContextPropagator {
+  override inject(
+    context: Context,
+    carrier: unknown,
+    setter: TextMapSetter,
+  ): void {
+    super.inject(context, carrier, {
+      set: (carrier, key, value) => {
+        if (key !== TRACE_STATE_KEY) {
+          setter.set(carrier, key, value)
+        }
+      },
+    })
+
+    const span = trace.getSpanContext(context)
+    if (span) {
+      const traceState = (span.traceState ?? createTraceState())
+        .set("sw", swValue(span))
+        .serialize()
+      setter.set(carrier, TRACE_STATE_KEY, traceState)
+    }
+  }
+}

packages/solarwinds-apm/test/propagation/trace-context.test.ts

@@ -0,0 +1,58 @@
+/*
+Copyright 2023-2024 SolarWinds Worldwide, LLC.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+import {
+  createTraceState,
+  defaultTextMapSetter,
+  ROOT_CONTEXT,
+  trace,
+  TraceFlags,
+} from "@opentelemetry/api"
+import { describe, expect, it } from "@solarwinds-apm/test"
+
+import { TraceContextPropagator } from "../../src/propagation/trace-context.js"
+
+describe("TraceContextPropagator", () => {
+  const propagator = new TraceContextPropagator()
+
+  it("inserts sw key in empty trace state", () => {
+    const headers: Partial<Record<string, string>> = {}
+    const context = trace.setSpanContext(ROOT_CONTEXT, {
+      spanId: "0123456789abcdef",
+      traceId: "0123456789abcdef0123456789abcdef",
+      traceFlags: TraceFlags.SAMPLED,
+    })
+
+    propagator.inject(context, headers, defaultTextMapSetter)
+    expect(headers.tracestate).to.equal("sw=0123456789abcdef-01")
+  })
+
+  it("inserts sw key in existing trace state", () => {
+    const headers: Partial<Record<string, string>> = {}
+    const context = trace.setSpanContext(ROOT_CONTEXT, {
+      spanId: "0123456789abcdef",
+      traceId: "0123456789abcdef0123456789abcdef",
+      traceFlags: TraceFlags.SAMPLED,
+      traceState: createTraceState().set("foo", "bar"),
+    })
+
+    propagator.inject(context, headers, defaultTextMapSetter)
+    expect(headers.tracestate).to.be.oneOf([
+      "sw=0123456789abcdef-01,foo=bar",
+      "foo=bar,sw=0123456789abcdef-01",
+    ])
+  })
+})