Support spans with manually specified trace and span ids

Author: ethercrow

Diff for: .gitignore

@@ -3,3 +3,9 @@
 /bash_completion.sh
 cabal.project.*
 dist-newstyle
+*.tracy
+*.eventlog
+*.trace.json
+*.gif
+*.webm
+*.log

Diff for: opentelemetry-extra/src/OpenTelemetry/EventlogStreaming_Internal.hs

@@ -74,12 +74,13 @@ data State = S
     threadMap :: IM.IntMap ThreadId,
     spanStacks :: HM.HashMap ThreadId (NonEmpty Span),
     traceMap :: HM.HashMap ThreadId TraceId,
+    specificSpans :: HM.HashMap Word64 Span,
     randomGen :: R.SMGen
   }
   deriving (Show)
 
 initialState :: Word64 -> R.SMGen -> State
-initialState timestamp = S timestamp mempty mempty mempty
+initialState timestamp = S timestamp mempty mempty mempty mempty
 
 processEvent :: Event -> State -> (State, [Span])
 processEvent (Event ts ev m_cap) st@(S {..}) =
@@ -102,6 +103,14 @@ processEvent (Event ts ev m_cap) st@(S {..}) =
         (HeapAllocated {allocBytes}, _, Just tid) ->
           (modifySpan tid (addEvent now "heap_alloc_bytes" (showT allocBytes)) st, [])
         (UserMessage {msg}, _, fromMaybe 1 -> tid) -> case T.words msg of
+          ("ot1" : "begin" : "specific" : "span" : trace_id_text : span_id_text : name) ->
+            let trace_id = TId (read ("0x" <> T.unpack trace_id_text))
+                span_id = SId (read ("0x" <> T.unpack span_id_text))
+             in beginSpecificSpan trace_id span_id (T.intercalate " " name) now st
+          ("ot1" : "end" : "specific" : "span" : trace_id_text : span_id_text : _) ->
+            let trace_id = TId (read ("0x" <> T.unpack trace_id_text))
+                span_id = SId (read ("0x" <> T.unpack span_id_text))
+             in endSpecificSpan span_id now st
           ("ot1" : "begin" : "span" : name) ->
             (pushSpan tid (T.intercalate " " name) now st, [])
           ("ot1" : "end" : "span" : _) -> popSpan tid now st
@@ -175,6 +184,44 @@ modifySpan tid f st =
         HM.update (\(sp :| sps) -> Just (f sp :| sps)) tid (spanStacks st)
     }
 
+beginSpecificSpan :: TraceId -> SpanId -> T.Text -> OTel.Timestamp -> State -> (State, [Span])
+beginSpecificSpan trace_id span_id@(SId s) name timestamp st =
+  case HM.lookup s (specificSpans st) of
+    Just sp -> (st {specificSpans = HM.delete s (specificSpans st)}, [sp {spanStartedAt = timestamp, spanOperation = name, spanContext = SpanContext span_id trace_id}])
+    Nothing ->
+      (st {specificSpans = HM.insert s sp (specificSpans st)}, [])
+      where
+        sp =
+          Span
+            { spanContext = SpanContext span_id trace_id,
+              spanOperation = name,
+              spanStartedAt = timestamp,
+              spanFinishedAt = 0,
+              spanTags = mempty,
+              spanEvents = mempty,
+              spanStatus = OK,
+              spanParentId = Nothing
+            }
+
+endSpecificSpan :: SpanId -> OTel.Timestamp -> State -> (State, [Span])
+endSpecificSpan span_id@(SId s) timestamp st =
+  case HM.lookup s (specificSpans st) of
+    Just sp -> (st {specificSpans = HM.delete s (specificSpans st)}, [sp {spanFinishedAt = timestamp}])
+    Nothing ->
+      (st {specificSpans = HM.insert s sp (specificSpans st)}, [])
+      where
+        sp =
+          Span
+            { spanContext = SpanContext span_id (TId 0),
+              spanOperation = "unknown",
+              spanStartedAt = 0,
+              spanFinishedAt = timestamp,
+              spanTags = mempty,
+              spanEvents = mempty,
+              spanStatus = OK,
+              spanParentId = Nothing
+            }
+
 pushSpan :: HasCallStack => ThreadId -> T.Text -> OTel.Timestamp -> State -> State
 pushSpan tid name timestamp st = st {spanStacks = new_stacks, randomGen = new_randomGen, traceMap = new_traceMap}
   where

Diff for: opentelemetry/src/OpenTelemetry/Eventlog.hs

@@ -52,3 +52,18 @@ withSpan operation action =
           liftIO endSpan
       )
       (\_span -> action)
+
+-- These two are supposed to be used when you have some custom control flow
+-- and a given span can begin on one thread and end on another. In this case
+-- the ordinary `beginSpan` and `endSpan` functions would assume a wrong thing
+-- and result in
+
+beginSpecificSpan :: TraceId -> SpanId -> String -> IO ()
+beginSpecificSpan (TId tid) (SId sid) k =
+  Debug.Trace.traceEventIO $
+    printf "ot1 begin specific span %d %d %s" tid sid k
+
+endSpecificSpan :: TraceId -> SpanId -> String -> IO ()
+endSpecificSpan (TId tid) (SId sid) k =
+  Debug.Trace.traceEventIO $
+    printf "ot1 end specific span %d %d %s" tid sid k