Merge pull request #17 from Yleisradio/fix-trace-context-propagation

Fix trace context propagation

Author: AriPaaWun

.github/workflows/build.yml

@@ -14,7 +14,7 @@ jobs:
     steps:
       - uses: actions/checkout@v4
 
-      - run: docker-compose up --detach
+      - run: docker compose up --detach
 
       - uses: actions/setup-java@v4
         with:

README.md

@@ -56,7 +56,8 @@ The options hash for `with-apm-transaction` accepts the following options:
 * `:labels` - `{String any}` - map or sequence of label names and values to add to the transaction  
 * `:tags` - `{String String}` - ~~map or sequence of label names and values to add to the transaction~~ Deprecated as of 0.5.0
 * `:activate?` - `Boolean` - whether to make the transaction active in the context of the executing thread (defaults to true). When activated, calling `apm/current-apm-transaction` returns this transaction.
-* `:traceparent` - `String` - the trace id when using APM's [Distributed tracing](https://www.elastic.co/guide/en/apm/get-started/current/distributed-tracing.html). Usually value is passed within HTTP request headers.
+* `:headers` - `{String String}` - a map of [trace context headers](https://www.w3.org/TR/trace-context/) (i.e. `traceparent` and `tracestate`) for [distributed tracing](https://www.elastic.co/guide/en/apm/get-started/current/distributed-tracing.html).
+* `:traceparent` - `String` - ~~the trace id when using APM's [Distributed tracing](https://www.elastic.co/guide/en/apm/get-started/current/distributed-tracing.html). Usually value is passed within HTTP request headers.~~ Deprecated as of 0.13.0
 
 The options hash for `with-apm-span` accepts the following options:
 

docker-compose.yml

@@ -1,7 +1,7 @@
 version: "3"
 services:
   elasticsearch:
-    image: "docker.elastic.co/elasticsearch/elasticsearch:7.16.0"
+    image: "docker.elastic.co/elasticsearch/elasticsearch:8.16.0"
     container_name: "elasticsearch"
     environment:
       - xpack.security.enabled=false
@@ -14,10 +14,10 @@ services:
       - "9200:9200"
       - "9300:9300"
   kibana:
-    image: "docker.elastic.co/kibana/kibana:7.16.0"
+    image: "docker.elastic.co/kibana/kibana:8.16.0"
     ports:
       - "5601:5601"
   apm-server:
-    image: "docker.elastic.co/apm/apm-server:7.16.0"
+    image: "docker.elastic.co/apm/apm-server:8.16.0"
     ports:
       - "8200:8200"

project.clj

@@ -1,4 +1,4 @@
-(defproject clojure-elastic-apm "0.12.0"
+(defproject clojure-elastic-apm "0.13.0"
   :description "Clojure wrapper for Elastic APM Java Agent"
   :url "https://github.com/Yleisradio/clojure-elastic-apm"
   :license {:name "Eclipse Public License"
@@ -9,6 +9,6 @@
                               "-Delastic.apm.application_packages=clojure-elastic-apm"
                               "-Delastic.apm.server_urls=http://localhost:8200"
                               "-Delastic.apm.metrics_interval=1s"]
-                   :dependencies [[clj-http "3.12.3"]
-                                  [cheshire "5.11.0"]]}
-             :provided {:dependencies [[co.elastic.apm/apm-agent-api "1.38.0"]]}})
+                   :dependencies [[clj-http "3.13.0"]
+                                  [cheshire "5.13.0"]]}
+             :provided {:dependencies [[co.elastic.apm/apm-agent-api "1.52.1"]]}})

src/clojure_elastic_apm/core.clj

@@ -1,5 +1,6 @@
 (ns clojure-elastic-apm.core
-  (:import [co.elastic.apm.api ElasticApm Transaction Span HeaderExtractor Outcome]))
+  (:require [clojure.string :as string])
+  (:import [co.elastic.apm.api ElasticApm Transaction Span HeaderExtractor HeadersExtractor Outcome]))
 
 (set! *warn-on-reflection* true)
 
@@ -13,10 +14,7 @@
   (.setLabel span-or-tx (name k) (str v)))
 
 (defn set-label [^Span span-or-tx k v]
-  (cond
-    (number? v) (.setLabel span-or-tx (name k) ^Number v)
-    (boolean? v) (.setLabel span-or-tx (name k) ^boolean v)
-    :else (.setLabel span-or-tx (name k) (str v))))
+  (.setLabel span-or-tx (name k) (str v)))
 
 (defn set-name [^Span span-or-tx name]
   (.setName span-or-tx name))
@@ -42,12 +40,24 @@
 (defn set-outcome-unknown [^Span span-or-tx]
   (set-outcome span-or-tx outcome-unknown))
 
-(defn trace-extractor ^HeaderExtractor [traceparent]
+(defn ^:deprecated trace-extractor
+  "DEPRECATED: Do not use. This function sets the value of every trace context
+  header (i.e. both \"traceparent\" and \"tracestate\") to the value you pass
+  it."
+  ^HeaderExtractor [traceparent]
   (reify HeaderExtractor
     (getFirstHeader [_ _] traceparent)))
 
 (def type-request Transaction/TYPE_REQUEST)
 
+(defn- header-values [headers header-name]
+  ;; "Where there are multiple headers with the same name, the adapter must
+  ;; concatenate the values into a single string, using the ASCII `,` character
+  ;; as a delimiter."
+  ;;
+  ;; -- https://github.com/ring-clojure/ring/blob/master/SPEC.md#headers
+  (some-> (get headers header-name) (string/split #",")))
+
 (defn start-transaction
   ([]
    (start-transaction {}))
@@ -56,9 +66,25 @@
      ^String tx-type :type
      tags :tags
      labels :labels
-     traceparent :traceparent}]
-   (let [tx (if traceparent
-              (ElasticApm/startTransactionWithRemoteParent (trace-extractor traceparent))
+     headers :headers
+     :as options}]
+   (let [tx (cond
+              headers
+              (ElasticApm/startTransactionWithRemoteParent
+                (reify HeaderExtractor
+                  (getFirstHeader [_ name]
+                    (first (header-values headers name))))
+                (reify HeadersExtractor
+                  (getAllHeaders [_ name]
+                    (header-values headers name))))
+
+              ;; Retained only for backwards compatibility. Use :headers
+              ;; instead.
+              (:traceparent options)
+              (ElasticApm/startTransactionWithRemoteParent
+                (trace-extractor (:traceparent options)))
+
+              :else
               (ElasticApm/startTransaction))]
      (when tx-name
        (set-name tx tx-name))

src/clojure_elastic_apm/ring.clj

@@ -39,21 +39,19 @@
    (fn
      ([{:keys [request-method uri headers] :as request}]
       (let [matched (match-patterns patterns uri)
-            tx-name (str (string/upper-case (name request-method)) " " matched)
-            traceparent (get headers "traceparent")]
+            tx-name (str (string/upper-case (name request-method)) " " matched)]
         (if matched
-          (with-apm-transaction [^Transaction tx {:name tx-name :type type-request :traceparent traceparent}]
+          (with-apm-transaction [^Transaction tx {:name tx-name :type type-request :headers headers}]
             (let [{:keys [status] :as response} (handler (assoc request :clojure-elastic-apm/transaction tx))]
               (when status
                 (.setResult tx (str "HTTP " status)))
               response))
           (handler request))))
      ([{:keys [request-method uri headers] :as request} respond raise]
       (let [matched (match-patterns patterns uri)
-            tx-name (str (string/upper-case (name request-method)) " " matched)
-            traceparent (get headers "traceparent")]
+            tx-name (str (string/upper-case (name request-method)) " " matched)]
         (if matched
-          (let [tx (apm/start-transaction {:name tx-name :type type-request :traceparent traceparent})
+          (let [tx (apm/start-transaction {:name tx-name :type type-request :headers headers})
                 req (assoc request :clojure-elastic-apm/transaction tx)]
             (with-open [_ (apm/activate tx)]
               (handler req

test/clojure_elastic_apm/core_test.clj

@@ -31,8 +31,8 @@
     (let [tx-details (es-find-first-document (str "(processor.event:transaction%20AND%20transaction.id:" @transaction-id ")"))]
       (is (= "TestTransaction" (get-in tx-details [:transaction :name])))
       (is (= "1" (get-in tx-details [:labels :t1])))
-      (is (= 2 (get-in tx-details [:labels :t2])))
-      (is (true? (get-in tx-details [:labels :t3])))
+      (is (= "2" (get-in tx-details [:labels :t2])))
+      (is (= "true" (get-in tx-details [:labels :t3])))
       (is (= "Label 4" (get-in tx-details [:labels :t4])))
       (is (= "test-result" (get-in tx-details [:transaction :result])))
       (is (= "success" (get-in tx-details [:event :outcome]))))))
@@ -159,3 +159,91 @@
              (-> doc
                  (select-keys [:span])
                  (update :span (fn [span] (select-keys span [:name :type])))))))))
+
+(defn trace-context
+  "Given an APM span, return the trace context[1] HTTP headers of the span.
+
+  [1]: https://www.w3.org/TR/trace-context/"
+  [span]
+  (let [headers (transient {})]
+    (.injectTraceHeaders span
+      (reify co.elastic.apm.api.HeaderInjector
+        (addHeader [_ name value]
+          (assoc! headers name value))))
+    (persistent! headers)))
+
+(def ^:private traceparent-regex
+  "https://www.w3.org/TR/trace-context/#traceparent-header"
+  #"^([0-9a-fA-F]{2})-([0-9a-fA-F]{32})-([0-9a-fA-F]{16})-([0-9a-fA-F]{2})$")
+
+(def ^:private tracestate-regex
+  "https://www.w3.org/TR/trace-context/#tracestate-header"
+  #"^[!-~]+(,[!-~]+=[!-~]+)*$")
+
+(defn parse-traceparent-header
+  "Given a traceparent header, parse the header value into its constituents."
+  [traceparent]
+  (zipmap [:version :trace-id :parent-id :trace-flags]
+    (rest (re-matches traceparent-regex traceparent))))
+
+(defn fields-except
+  "Given a traceparent header, return a map of its constituents, except fields."
+  [traceparent & fields]
+  (-> (parse-traceparent-header traceparent)
+      (apply dissoc fields)))
+
+(deftest trace-context-test
+  (testing "no opts"
+    (let [tx (apm/start-transaction)
+          {:strs [traceparent tracestate]} (trace-context tx)]
+      ;; APM adds new trace context headers.
+      (is (re-matches traceparent-regex traceparent))
+      (is (re-matches tracestate-regex tracestate))))
+
+  (testing ":headers -- empty & nil"
+    (let [tx (apm/start-transaction {:headers {}})
+          {:strs [traceparent tracestate]} (trace-context tx)]
+      (is (re-matches traceparent-regex traceparent))
+      (is (re-matches tracestate-regex tracestate)))
+
+    (let [tx (apm/start-transaction {:headers nil})
+          {:strs [traceparent tracestate]} (trace-context tx)]
+      (is (re-matches traceparent-regex traceparent))
+      (is (re-matches tracestate-regex tracestate))))
+
+  (testing ":headers -- traceparent only"
+    (let [inbound-traceparent "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"
+          tx (apm/start-transaction {:headers {"traceparent" inbound-traceparent}})
+          {:strs [traceparent]} (trace-context tx)]
+      (is (= (fields-except inbound-traceparent :parent-id)
+             (fields-except traceparent :parent-id)))))
+
+  (testing ":headers -- traceparent & tracestate"
+    (let [inbound-traceparent "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"
+          inbound-tracestate "es=s:0.25"
+          tx (apm/start-transaction {:headers {"traceparent" inbound-traceparent
+                                               "tracestate" inbound-tracestate}})
+          {:strs [traceparent tracestate]} (trace-context tx)]
+      (is (= (fields-except inbound-traceparent :parent-id)
+             (fields-except traceparent :parent-id)))
+      (is (= inbound-tracestate tracestate))))
+
+  (testing ":headers -- multi-value tracestate"
+    (let [inbound-traceparent "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"
+          inbound-tracestate "es=s:0.5,rojo=00f067aa0ba902b7"
+          tx (apm/start-transaction {:headers {"traceparent" inbound-traceparent
+                                               "tracestate" inbound-tracestate}})
+          {:strs [tracestate]} (trace-context tx)]
+      (is (= inbound-tracestate tracestate))))
+
+  (testing ":traceparent (backwards compatibility only, do not use)"
+    (let [inbound-traceparent "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"
+          tx (apm/start-transaction {:traceparent inbound-traceparent})
+          {:strs [traceparent tracestate]} (trace-context tx)]
+      (is (= (fields-except inbound-traceparent :parent-id)
+             (fields-except traceparent :parent-id)))
+
+      ;; Using the :traceparent option incorrectly sets the value of every
+      ;; trace context header -- including tracestate -- to the value of the
+      ;; traceparent header. Do not use it.
+      (is (= inbound-traceparent tracestate)))))