Add trace internal capabilities

.cljstyle

@@ -66,6 +66,7 @@
     smart-map                          [[:inner 0]]
     some                               [[:inner 0]]
     swap!                              [[:inner 0]]
+    t/with-span!                       [[:inner 0]]
     update                             [[:inner 0]]
     update-in                          [[:inner 0]]}}
 

resources/clj-kondo.exports/com.wsscode/pathom3/config.edn

@@ -1,5 +1,6 @@
 {:lint-as {com.wsscode.pathom3.connect.operation/defmutation clojure.core/defn
            com.wsscode.pathom3.connect.operation/defresolver clojure.core/defn
            com.wsscode.pathom3.plugin/defplugin              clojure.core/def
+           com.wsscode.pathom3.trace/with-span!              clojure.core/let
            com.wsscode.promesa.macros/clet                   clojure.core/let
            com.wsscode.promesa.macros/ctry                   clojure.core/try}}

src/main/com/wsscode/pathom3/trace.cljc

@@ -0,0 +1,156 @@
+(ns com.wsscode.pathom3.trace
+  (:require
+    [clojure.set :as set]
+    [clojure.spec.alpha :as s]
+    [com.fulcrologic.guardrails.core :refer [>def]]
+    [com.wsscode.misc.coll :as coll]
+    [com.wsscode.misc.refs :as refs]
+    [com.wsscode.misc.time :as time])
+  #?(:clj
+     (:import
+       (java.security
+         SecureRandom))))
+
+; region specs
+
+(>def ::type #{::type-close-span
+               ::type-log
+               ::type-set-attributes})
+
+(>def ::attributes (s/keys))
+(>def ::span-id string?)
+
+(>def ::timestamp pos-int?)
+(>def ::start-time ::timestamp)
+(>def ::end-time ::timestamp)
+
+(>def ::events (s/coll-of (s/keys)))
+
+(>def ::name string?)
+(>def ::error string?)
+
+(>def ::parent-span-id ::span-id)
+
+(>def ::trace* "Atom holding trace data" refs/atom?)
+
+; endregion
+
+; region id generation
+
+#?(:clj
+   (defn generate-bytes
+     "Generates a 64-bit Span ID for OpenTelemetry."
+     [size]
+     (let [random-bytes (byte-array size)]
+       (.nextBytes (SecureRandom.) random-bytes)
+       random-bytes)))
+
+#?(:clj
+   (defn bytes-to-hex
+     "Converts a byte array to a hex string."
+     [byte-array]
+     (apply str "0x" (mapv #(format "%02x" (bit-and 0xFF %)) byte-array))))
+
+(defn new-trace-id []
+  #?(:clj  (bytes-to-hex (generate-bytes 16))
+     :cljs (gensym "trace-id")))
+
+(defn new-span-id []
+  #?(:clj  (bytes-to-hex (generate-bytes 8))
+     :cljs (gensym "span-id")))
+
+; endregion
+
+(defn add-signal!
+  "Adds a signal to the trace. This is a low-level function, you should use the other functions to add signals to the trace."
+  [{::keys [trace* parent-span-id]} entry]
+  (when trace*
+    (swap! trace* conj (assoc entry ::timestamp (time/now-ms)
+                         ::parent-span-id parent-span-id))))
+
+(defn open-span!
+  "Opens a new span and adds it to the trace. Returns the span id."
+  [env span]
+  (let [id (new-span-id)]
+    (add-signal! env (-> span
+                         (assoc ::span-id id)
+                         (update ::attributes assoc :com.wsscode.pathom3.path/path (get env :com.wsscode.pathom3.path/path []))))
+    id))
+
+(defn close-span!
+  "Closes a span and adds it to the trace."
+  [env span-id]
+  (add-signal! env {::type ::type-close-span ::close-span-id span-id}))
+
+(defn under-span
+  "Returns a new environment setting the context span id."
+  [env span-id]
+  (assoc env ::parent-span-id span-id))
+
+#?(:clj
+   (defmacro with-span!
+     "Opens a new span and closes it after the body is executed. The span id is bound to the environment.
+
+        (t/with-span! [env {::t/env env}]
+          (do-something))"
+     [[sym span] & body]
+     `(if-let [env# (get ~span ::env)]
+        (let [span#    (dissoc ~span ::env)
+              span-id# (open-span! env# span#)
+              res#     (let [~sym (under-span env# span-id#)]
+                         ~@body)]
+          (close-span! env# span-id#)
+          res#)
+        (throw (ex-info "With span requires environment as part of the data" {})))))
+
+#?(:clj
+   (s/fdef with-span!
+     :args (s/cat :binding (s/tuple simple-symbol? any?) :body any?)
+     :ret any?))
+
+(defn log!
+  "Adds a log entry to the trace. This will end up as an event to the current span."
+  [env log-entry]
+  (add-signal! env (assoc log-entry ::type ::type-log)))
+
+(defn set-attributes! [env & {:as attributes}]
+  (add-signal! env {::type       ::type-set-attributes
+                    ::attributes attributes}))
+
+(defn fold-trace
+  "Folds the trace data into spans. This materializes the trace data, making it suitable for usage."
+  [trace]
+  (let [{:keys [top-level props]}
+        (-> (reduce
+              (fn [out {::keys [type] :as entry}]
+                (case type
+                  (::type-log ::type-set-attributes)
+                  (update-in out [:props (::parent-span-id entry)] coll/vconj entry)
+
+                  ::type-close-span
+                  (update-in out [:props (::close-span-id entry)] coll/vconj entry)
+
+                  ; else
+                  (update out :top-level conj! (set/rename-keys entry {::timestamp ::start-time}))))
+              {:top-level (transient [])
+               :props     {}}
+              trace)
+            (update :top-level persistent!))]
+    (->> top-level
+         (map
+           (fn [entry]
+             (if-let [props (get props (::span-id entry))]
+               (reduce
+                 (fn [entry {::keys [type] :as ext}]
+                   (case type
+                     ::type-log
+                     (update entry ::events coll/vconj (dissoc ext ::parent-span-id))
+
+                     ::type-set-attributes
+                     (update entry ::attributes merge (::attributes ext))
+
+                     ::type-close-span
+                     (assoc entry ::end-time (::timestamp ext))))
+                 entry
+                 props)
+               entry))))))

test/com/wsscode/pathom3/trace_test.cljc

@@ -0,0 +1,162 @@
+(ns com.wsscode.pathom3.trace-test
+  (:require
+    [clojure.test :refer [deftest is testing]]
+    [com.wsscode.misc.time :as time]
+    [com.wsscode.pathom3.trace :as t]))
+
+(deftest add-signal!-test
+  (testing "adds when trace* is present"
+    (with-redefs [time/now-ms (constantly 123)]
+      (let [trace* (atom [])
+            env    {::t/trace* trace*}]
+        (t/add-signal! env {::t/log "test"})
+        (is (= @trace* [{::t/log            "test"
+                         ::t/timestamp      123
+                         ::t/parent-span-id nil}])))))
+
+  (testing "does not add when trace* is not present"
+    (with-redefs [time/now-ms (constantly 123)]
+      (let [trace* (atom [])
+            env    {}]
+        (t/add-signal! env {::t/log "test"})
+        (is (= @trace* []))))))
+
+(deftest open-span!-test
+  (testing "opens a span and adds it to the trace"
+    (with-redefs [t/new-span-id (constantly "span-id")
+                  time/now-ms   (constantly 123)]
+      (let [trace* (atom [])
+            env    {::t/trace* trace*}]
+        (is (= (t/open-span! env {::t/log "test"})
+               "span-id"))
+        (is (= @trace* [{::t/log            "test"
+                         ::t/timestamp      123
+                         ::t/parent-span-id nil
+                         ::t/attributes     {:com.wsscode.pathom3.path/path []}
+                         ::t/span-id        "span-id"}]))))))
+
+(deftest close-span!-test
+  (testing "closes a span and adds it to the trace"
+    (with-redefs [t/new-span-id (constantly "span-id")
+                  time/now-ms   (constantly 123)]
+      (let [trace* (atom [])
+            env    {::t/trace* trace*}]
+        (t/close-span! env "span-id")
+        (is (= @trace* [{:com.wsscode.pathom3.trace/close-span-id  "span-id"
+                         :com.wsscode.pathom3.trace/parent-span-id nil
+                         :com.wsscode.pathom3.trace/timestamp      123
+                         :com.wsscode.pathom3.trace/type           :com.wsscode.pathom3.trace/type-close-span}]))))))
+
+(deftest under-span-test
+  (testing "returns a new environment setting the context span id"
+    (let [env {::t/parent-span-id "span-id"}]
+      (is (= (t/under-span env "span-id-2")
+             {::t/parent-span-id "span-id-2"})))))
+
+(deftest with-span!-test
+  (testing "opens a new span and closes it after the body is executed"
+    (with-redefs [t/new-span-id (constantly "span-id")
+                  time/now-ms   (constantly 123)]
+      (let [trace* (atom [])
+            env    {::t/trace* trace*}]
+        (t/with-span! [env {::t/env env}]
+          (t/log! env {::t/name "test"}))
+        (is (= @trace*
+               [{:com.wsscode.pathom3.trace/attributes     {:com.wsscode.pathom3.path/path []}
+                 :com.wsscode.pathom3.trace/parent-span-id nil
+                 :com.wsscode.pathom3.trace/span-id        "span-id"
+                 :com.wsscode.pathom3.trace/timestamp      123}
+                {:com.wsscode.pathom3.trace/name           "test"
+                 :com.wsscode.pathom3.trace/parent-span-id "span-id"
+                 :com.wsscode.pathom3.trace/timestamp      123
+                 :com.wsscode.pathom3.trace/type           :com.wsscode.pathom3.trace/type-log}
+                {:com.wsscode.pathom3.trace/close-span-id  "span-id"
+                 :com.wsscode.pathom3.trace/parent-span-id nil
+                 :com.wsscode.pathom3.trace/timestamp      123
+                 :com.wsscode.pathom3.trace/type           :com.wsscode.pathom3.trace/type-close-span}]))))))
+
+(deftest log!-test
+  (testing "adds a log entry"
+    (with-redefs [t/new-span-id (constantly "span-id")
+                  time/now-ms   (constantly 123)]
+      (let [trace* (atom [])
+            env    {::t/trace* trace*}]
+        (t/log! env {::t/name "test"})
+        (is (= @trace*
+               [{:com.wsscode.pathom3.trace/name           "test"
+                 :com.wsscode.pathom3.trace/parent-span-id nil
+                 :com.wsscode.pathom3.trace/timestamp      123
+                 :com.wsscode.pathom3.trace/type           :com.wsscode.pathom3.trace/type-log}]))))))
+
+(deftest set-attributes!-test
+  (testing "adds an attribute set entry"
+    (with-redefs [t/new-span-id (constantly "span-id")
+                  time/now-ms   (constantly 123)]
+      (let [trace* (atom [])
+            env    {::t/trace* trace*}]
+        (t/set-attributes! env :some-attr "value")
+        (is (= @trace*
+               [{:com.wsscode.pathom3.trace/attributes     {:some-attr "value"}
+                 :com.wsscode.pathom3.trace/parent-span-id nil
+                 :com.wsscode.pathom3.trace/timestamp      123
+                 :com.wsscode.pathom3.trace/type           :com.wsscode.pathom3.trace/type-set-attributes}]))))))
+
+(deftest fold-trace-test
+  (testing "closes spans"
+    (is (= (t/fold-trace
+             [{:com.wsscode.pathom3.trace/attributes     {:com.wsscode.pathom3.path/path []}
+               :com.wsscode.pathom3.trace/parent-span-id nil
+               :com.wsscode.pathom3.trace/span-id        "span-id"
+               :com.wsscode.pathom3.trace/timestamp      123}
+              {:com.wsscode.pathom3.trace/close-span-id "span-id"
+               :com.wsscode.pathom3.trace/timestamp     1234
+               :com.wsscode.pathom3.trace/type          :com.wsscode.pathom3.trace/type-close-span}])
+           [{:com.wsscode.pathom3.trace/attributes     {:com.wsscode.pathom3.path/path []},
+             :com.wsscode.pathom3.trace/parent-span-id nil,
+             :com.wsscode.pathom3.trace/span-id        "span-id",
+             :com.wsscode.pathom3.trace/start-time     123,
+             :com.wsscode.pathom3.trace/end-time       1234}])))
+
+  (testing "add log to span"
+    (is (= (t/fold-trace
+             [{:com.wsscode.pathom3.trace/attributes     {:com.wsscode.pathom3.path/path []}
+               :com.wsscode.pathom3.trace/parent-span-id nil
+               :com.wsscode.pathom3.trace/span-id        "span-id"
+               :com.wsscode.pathom3.trace/timestamp      123}
+              {:com.wsscode.pathom3.trace/attributes     {:com.wsscode.pathom3.path/path []}
+               :com.wsscode.pathom3.trace/name           "test"
+               :com.wsscode.pathom3.trace/parent-span-id "span-id"
+               :com.wsscode.pathom3.trace/timestamp      124
+               :com.wsscode.pathom3.trace/type           :com.wsscode.pathom3.trace/type-log}
+              {:com.wsscode.pathom3.trace/close-span-id "span-id"
+               :com.wsscode.pathom3.trace/timestamp     125
+               :com.wsscode.pathom3.trace/type          :com.wsscode.pathom3.trace/type-close-span}])
+           [{:com.wsscode.pathom3.trace/attributes     {:com.wsscode.pathom3.path/path []},
+             :com.wsscode.pathom3.trace/parent-span-id nil,
+             :com.wsscode.pathom3.trace/span-id        "span-id",
+             :com.wsscode.pathom3.trace/start-time     123,
+             :com.wsscode.pathom3.trace/events         [{:com.wsscode.pathom3.trace/attributes {:com.wsscode.pathom3.path/path []},
+                                                         :com.wsscode.pathom3.trace/name       "test",
+                                                         :com.wsscode.pathom3.trace/timestamp  124,
+                                                         :com.wsscode.pathom3.trace/type       :com.wsscode.pathom3.trace/type-log}],
+             :com.wsscode.pathom3.trace/end-time       125}])))
+
+  (testing "set attributes of span"
+    (is (= (t/fold-trace
+             [{:com.wsscode.pathom3.trace/attributes     {:com.wsscode.pathom3.path/path []}
+               :com.wsscode.pathom3.trace/parent-span-id nil
+               :com.wsscode.pathom3.trace/span-id        "span-id"
+               :com.wsscode.pathom3.trace/timestamp      123}
+              {:com.wsscode.pathom3.trace/attributes     {:some-attr "value"}
+               :com.wsscode.pathom3.trace/parent-span-id "span-id"
+               :com.wsscode.pathom3.trace/timestamp      123
+               :com.wsscode.pathom3.trace/type           :com.wsscode.pathom3.trace/type-set-attributes}
+              {:com.wsscode.pathom3.trace/close-span-id "span-id"
+               :com.wsscode.pathom3.trace/timestamp     125
+               :com.wsscode.pathom3.trace/type          :com.wsscode.pathom3.trace/type-close-span}])
+           [{:com.wsscode.pathom3.trace/attributes     {:com.wsscode.pathom3.path/path []
+                                                        :some-attr                     "value"}
+             :com.wsscode.pathom3.trace/end-time       125
+             :com.wsscode.pathom3.trace/parent-span-id nil
+             :com.wsscode.pathom3.trace/span-id        "span-id"
+             :com.wsscode.pathom3.trace/start-time     123}]))))