thread-protected automatic trace id propagation to child threads

10 files changed, 89 insertions, 4 deletions

diff --git a/compiler/GHC/Builtin/primops.txt.pp b/compiler/GHC/Builtin/primops.txt.pp
index 1e3b9b8af5..cd4a337dde 100644
--- a/compiler/GHC/Builtin/primops.txt.pp
+++ b/compiler/GHC/Builtin/primops.txt.pp
@@ -2935,6 +2935,18 @@ primop  MyThreadIdOp "myThreadId#" GenPrimOp
    with
    has_side_effects = True
 
+primop SetAdamTraceIdOp "setAdamTraceId#" GenPrimOp
+   ByteArray# -> State# s -> State# s
+   with
+   has_side_effects = True
+   out_of_line      = True
+
+primop GetAdamTraceIdOp "getAdamTraceId#" GenPrimOp
+   State# s -> (# State# s, ByteArray# #)
+   with
+   has_side_effects = True
+   out_of_line      = True
+
 primop LabelThreadOp "labelThread#" GenPrimOp
    ThreadId# -> Addr# -> State# RealWorld -> State# RealWorld
    with
diff --git a/compiler/GHC/StgToCmm/Prim.hs b/compiler/GHC/StgToCmm/Prim.hs
index b0f9fddad6..8e95c50a2f 100644
--- a/compiler/GHC/StgToCmm/Prim.hs
+++ b/compiler/GHC/StgToCmm/Prim.hs
@@ -1485,6 +1485,8 @@ emitPrimOp dflags = \case
   KillThreadOp -> alwaysExternal
   YieldOp -> alwaysExternal
   LabelThreadOp -> alwaysExternal
+  GetAdamTraceIdOp -> alwaysExternal
+  SetAdamTraceIdOp -> alwaysExternal
   IsCurrentThreadBoundOp -> alwaysExternal
   NoDuplicateOp -> alwaysExternal
   ThreadStatusOp -> alwaysExternal
diff --git a/includes/rts/storage/TSO.h b/includes/rts/storage/TSO.h
index 3a488d97b5..47fa0501d5 100644
--- a/includes/rts/storage/TSO.h
+++ b/includes/rts/storage/TSO.h
@@ -182,6 +182,17 @@ typedef struct StgTSO_ {
 #if defined(mingw32_HOST_OS)
     StgWord32 saved_winerror;
 #endif
+    /** For automatic trace id propagation (OpenTelemetry) across threads
+     *  and futher to external processes. GHC Eventlog allows to persist trace id
+     *  at runtime and reconstruct it only at parsing logs.
+
+     *  setAdamTraceId trId
+     *  forkIO $ forkIO $ forkIO $ doHttpCall (TraceHeader <$> getAdamTraceId)
+     *   or
+     *  setAdamTraceId trId
+     *  forkIO $ $forkIO $ forkIO $ runProcess $ (EnvVariable TraceId) <$> getAdamTraceId
+     */
+    StgClosure* trace_id;
 
 } *StgTSOPtr; // StgTSO defined in rts/Types.h
 
diff --git a/includes/stg/MiscClosures.h b/includes/stg/MiscClosures.h
index dc2b0715ca..ac27a4e5c9 100644
--- a/includes/stg/MiscClosures.h
+++ b/includes/stg/MiscClosures.h
@@ -449,6 +449,8 @@ RTS_FUN_DECL(stg_maskUninterruptiblezh);
 RTS_FUN_DECL(stg_unmaskAsyncExceptionszh);
 RTS_FUN_DECL(stg_myThreadIdzh);
 RTS_FUN_DECL(stg_labelThreadzh);
+RTS_FUN_DECL(stg_setAdamTraceIdzh);
+RTS_FUN_DECL(stg_getAdamTraceIdzh);
 RTS_FUN_DECL(stg_isCurrentThreadBoundzh);
 RTS_FUN_DECL(stg_threadStatuszh);
 
diff --git a/libraries/base/GHC/Conc/Sync.hs b/libraries/base/GHC/Conc/Sync.hs
index d6ffbc2de9..c1cd6a648f 100644
--- a/libraries/base/GHC/Conc/Sync.hs
+++ b/libraries/base/GHC/Conc/Sync.hs
@@ -33,7 +33,7 @@
 -- #not-home
 module GHC.Conc.Sync
         ( ThreadId(..)
-
+        , AdamTraceId(..)
         -- * Forking and suchlike
         , forkIO
         , forkIOWithUnmask
@@ -46,6 +46,8 @@ module GHC.Conc.Sync
         , numSparks
         , childHandler
         , myThreadId
+        , setAdamTraceId
+        , getAdamTraceId
         , killThread
         , throwTo
         , par
@@ -116,6 +118,7 @@ import GHC.Real         ( fromIntegral )
 import GHC.Show         ( Show(..), showParen, showString )
 import GHC.Stable       ( StablePtr(..) )
 import GHC.Weak
+import GHC.Word
 
 import Unsafe.Coerce    ( unsafeCoerce# )
 
@@ -151,6 +154,8 @@ instance Show ThreadId where
         showString "ThreadId " .
         showsPrec d (getThreadId (id2TSO t))
 
+data AdamTraceId = AdamTraceId ByteArray#
+
 foreign import ccall unsafe "rts_getThreadId" getThreadId :: ThreadId# -> CInt
 
 id2TSO :: ThreadId -> ThreadId#
@@ -471,6 +476,17 @@ myThreadId = IO $ \s ->
    case (myThreadId# s) of (# s1, tid #) -> (# s1, ThreadId tid #)
 
 
+-- | Returns the value associated with origin thread (GHC only).
+getAdamTraceId :: IO AdamTraceId
+getAdamTraceId = IO $ \s ->
+   case (getAdamTraceId# s) of (# s1, trId #) -> (# s1, AdamTraceId trId #)
+
+-- | Associates value with current thread which is transfered
+-- to all child threads (GHC only).
+setAdamTraceId :: AdamTraceId -> IO ()
+setAdamTraceId (AdamTraceId trId) = IO $ \s ->
+   case (setAdamTraceId# trId s) of s1 -> (# s1, () #)
+
 -- | The 'yield' action allows (forces, in a co-operative multitasking
 -- implementation) a context-switch to any other currently runnable
 -- threads (if any), and is occasionally useful when implementing
diff --git a/rts/PrimOps.cmm b/rts/PrimOps.cmm
index 048cde8065..d4c9bbba8d 100644
--- a/rts/PrimOps.cmm
+++ b/rts/PrimOps.cmm
@@ -1042,6 +1042,19 @@ stg_yieldzh ()
     jump stg_yield_noregs();
 }
 
+stg_setAdamTraceIdzh ( gcptr src )
+{
+    ccall c_setAdamTraceId(CurrentTSO, src);
+    return ();
+}
+
+stg_getAdamTraceIdzh ( /* no args */ )
+{
+    gcptr ret;
+    (ret) = ccall c_getAdamTraceId(CurrentTSO);
+    return (ret);
+}
+
 stg_labelThreadzh ( gcptr threadid, W_ addr )
 {
 #if defined(DEBUG) || defined(TRACING) || defined(DTRACE)
diff --git a/rts/RtsSymbols.c b/rts/RtsSymbols.c
index 7e89aaffc5..90940f1a11 100644
--- a/rts/RtsSymbols.c
+++ b/rts/RtsSymbols.c
@@ -681,6 +681,8 @@
       SymI_HasProto(stg_makeStablePtrzh)                                \
       SymI_HasProto(stg_mkApUpd0zh)                                     \
       SymI_HasProto(stg_labelThreadzh)                                  \
+      SymI_HasProto(stg_setAdamTraceIdzh)                               \
+      SymI_HasProto(stg_getAdamTraceIdzh)                               \
       SymI_HasProto(stg_newArrayzh)                                     \
       SymI_HasProto(stg_copyArrayzh)                                    \
       SymI_HasProto(stg_copyMutableArrayzh)                             \
diff --git a/rts/Threads.c b/rts/Threads.c
index 22d58bb48b..e19585fbd6 100644
--- a/rts/Threads.c
+++ b/rts/Threads.c
@@ -112,6 +112,7 @@ createThread(Capability *cap, W_ size)
     ASSIGN_Int64((W_*)&(tso->alloc_limit), 0);
 
     tso->trec = NO_TREC;
+    tso->trace_id = 0;
 
 #if defined(PROFILING)
     tso->prof.cccs = CCS_MAIN;
@@ -132,12 +133,37 @@ createThread(Capability *cap, W_ size)
     g0->threads = tso;
     RELEASE_LOCK(&sched_mutex);
 
+    if (cap->r.rCurrentTSO) {
+        tso->trace_id = cap->r.rCurrentTSO->trace_id;
+    }
+
     // ToDo: report the stack size in the event?
     traceEventCreateThread(cap, tso);
 
     return tso;
 }
 
+void c_setAdamTraceId(StgTSO* tso USED_IF_THREADS, StgClosure* trId) {
+    tso->trace_id = trId;
+    write_barrier();
+    dirty_TSO(tso->cap, tso);
+}
+
+#define SIZEOF_W  sizeof(void*)
+#define ROUNDUP_BYTES_TO_WDS(n) (((n) + SIZEOF_W - 1) / SIZEOF_W)
+#define SIZEOF_StgArrBytes sizeof(StgArrBytes)
+
+StgClosure* c_getAdamTraceId(StgTSO* tso USED_IF_THREADS) {
+    if (tso->trace_id) {
+        return tso->trace_id;
+    }
+    StgArrBytes* trace_id = (StgArrBytes*)allocate(tso->cap, ROUNDUP_BYTES_TO_WDS(SIZEOF_StgArrBytes));
+    SET_HDR(trace_id, &stg_ARR_WORDS_info, CCS_SYSTEM);
+    trace_id->bytes = 0;
+    return trace_id;
+}
+
+
 /* ---------------------------------------------------------------------------
  * Equality on Thread ids.
  *
@@ -902,7 +928,6 @@ printThreadBlockage(StgTSO *tso)
   }
 }
 
-
 void
 printThreadStatus(StgTSO *t)
 {
diff --git a/rts/sm/Scav.c b/rts/sm/Scav.c
index 501d958aae..b89554cdc4 100644
--- a/rts/sm/Scav.c
+++ b/rts/sm/Scav.c
@@ -123,7 +123,9 @@ scavengeTSO (StgTSO *tso)
     evacuate((StgClosure **)&tso->trec);
 
     evacuate((StgClosure **)&tso->stackobj);
-
+    if (tso->trace_id) {
+        evacuate((StgClosure **)&tso->trace_id);
+    }
     evacuate((StgClosure **)&tso->_link);
     if (   tso->why_blocked == BlockedOnMVar
         || tso->why_blocked == BlockedOnMVarRead
diff --git a/utils/deriveConstants/Main.hs b/utils/deriveConstants/Main.hs
index 1867d824b6..06c5ca79b9 100644
--- a/utils/deriveConstants/Main.hs
+++ b/utils/deriveConstants/Main.hs
@@ -443,6 +443,7 @@ wanteds os = concat
           ,closureField  C    "StgTSO"      "what_next"
           ,closureField  C    "StgTSO"      "why_blocked"
           ,closureField  C    "StgTSO"      "block_info"
+          ,closureField  C    "StgTSO"      "trace_id"
           ,closureField  C    "StgTSO"      "blocked_exceptions"
           ,closureField  C    "StgTSO"      "id"
           ,closureField  C    "StgTSO"      "cap"
@@ -969,4 +970,3 @@ execute verbose prog args
       ec <- rawSystem prog args
       unless (ec == ExitSuccess) $
           die ("Executing " ++ show prog ++ " failed")
-