HPCC-31425 Add support for tracking the active span in the thread con…

…text

Signed-off-by: Gavin Halliday <gavin.halliday@lexisnexis.com>

system/jlib/jthread.cpp

@@ -2684,12 +2684,14 @@ static TraceFlags defaultTraceFlags = TraceFlags::Standard;
 static thread_local LogMsgJobId defaultJobId = UnknownJob;
 static thread_local TraceFlags threadTraceFlags = TraceFlags::Standard;
 static thread_local const IContextLogger *default_thread_logctx = nullptr;
+static thread_local ISpan * threadActiveSpan = nullptr;
 
 void saveThreadContext(SavedThreadContext & saveCtx)
 {
     saveCtx.jobId = defaultJobId;
     saveCtx.logctx = default_thread_logctx;
     saveCtx.traceFlags = threadTraceFlags;
+    saveCtx.activeSpan = threadActiveSpan;
 }
 
 void restoreThreadContext(const SavedThreadContext & saveCtx)
@@ -2700,9 +2702,9 @@ void restoreThreadContext(const SavedThreadContext & saveCtx)
     defaultJobId = saveCtx.jobId;
     default_thread_logctx = saveCtx.logctx;
     threadTraceFlags = saveCtx.traceFlags;
+    threadActiveSpan = saveCtx.activeSpan;
 }
 
-
 LogMsgJobId queryThreadedJobId()
 {
     return defaultJobId;
@@ -2718,6 +2720,21 @@ const IContextLogger * queryThreadedContextLogger()
     return default_thread_logctx;
 }
 
+ISpan * queryThreadedActiveSpan()
+{
+    ISpan * result = threadActiveSpan;
+    if (!result)
+        result = queryNullSpan();
+    return result;
+}
+
+ISpan * setThreadedActiveSpan(ISpan * span)
+{
+    ISpan * ret = threadActiveSpan;
+    threadActiveSpan = span;
+    return ret;
+}
+
 //---------------------------
 
 bool doTrace(TraceFlags featureFlag, TraceFlags level)

system/jlib/jthread.hpp

@@ -34,18 +34,27 @@
 #endif
 
 // Functions used to reset thread-local context variables, when a threadpool starts
+//NOTE: Currently activeSpan is not linked in the SavedThreadContext.
+//That will not cause any problems as long as the lifetime of the thread executing with that saved span value is less than the lifetime of the span.
+//It would be trivial to use a linked pointer inside the savedthread context, at the expense of a little inefficiency.
+//Revisit when this feature starts being used in anger.
 typedef unsigned __int64 LogMsgJobId;
 interface IContextLogger;
+interface ISpan;
 struct jlib_decl SavedThreadContext
 {
     const IContextLogger * logctx = nullptr;
     LogMsgJobId jobId = (LogMsgJobId)-1;
     TraceFlags traceFlags = queryDefaultTraceFlags();
+    ISpan * activeSpan = nullptr;
 };
 
 extern void restoreThreadContext(const SavedThreadContext & saveCtx);
 extern void saveThreadContext(SavedThreadContext & saveCtx);
 
+extern jlib_decl ISpan * queryThreadedActiveSpan();
+extern jlib_decl ISpan * setThreadedActiveSpan(ISpan * span);
+
 //--------------------------------------------------------------
 
 interface jlib_decl IThread : public IInterface

system/jlib/jtrace.cpp

@@ -1458,6 +1458,43 @@ ISpan * CTraceManager::createServerSpan(const char * name, const IProperties * h
 
 //---------------------------------------------------------------------------------------------------------------------
 
+OwnedSpanScope::OwnedSpanScope(ISpan * _ptr) : span(_ptr)
+{
+    if (_ptr)
+        prevSpan = setThreadedActiveSpan(_ptr);
+}
+
+void OwnedSpanScope::setown(ISpan * _span)
+{
+    assertex(_span);
+    //Just in case the span is already set, ensure it is ended and that the previous span is restored.
+    clear();
+    span.setown(_span);
+    prevSpan = setThreadedActiveSpan(_span);
+}
+
+void OwnedSpanScope::set(ISpan * _span)
+{
+    setown(LINK(_span));
+}
+
+void OwnedSpanScope::clear()
+{
+    if (span)
+    {
+        setThreadedActiveSpan(prevSpan);
+        span->endSpan();
+        span.clear();
+    }
+}
+
+OwnedSpanScope::~OwnedSpanScope()
+{
+    clear();
+}
+
+//---------------------------------------------------------------------------------------------------------------------
+
 MODULE_INIT(INIT_PRIORITY_STANDARD)
 {
     return true;
@@ -1475,6 +1512,11 @@ ISpan * getNullSpan()
     return nullSpan.getLink();
 }
 
+ISpan * queryNullSpan()
+{
+    return nullSpan;
+}
+
 void initTraceManager(const char * componentName, const IPropertyTree * componentConfig, const IPropertyTree * globalConfig)
 {
     theTraceManager.query([=] () { return new CTraceManager(componentName, componentConfig, globalConfig); });

system/jlib/jtrace.hpp

@@ -134,16 +134,24 @@ interface ISpan : extends IInterface
     virtual const char* queryLocalId() const = 0;
 };
 
-class OwnedSpanScope : public Owned<ISpan>
+class jlib_decl OwnedSpanScope
 {
 public:
-    inline OwnedSpanScope() { }
-    inline OwnedSpanScope(ISpan * _ptr) : Owned<ISpan>(_ptr) { }
-    ~OwnedSpanScope()
-    {
-        if (get())
-            get()->endSpan();
-    }
+    OwnedSpanScope() = default;
+    OwnedSpanScope(ISpan * _ptr);
+    ~OwnedSpanScope();
+
+    inline ISpan * operator -> () const         { return span; }
+    inline operator ISpan *() const             { return span; }
+
+    void clear();
+    ISpan * query() const { return span; }
+    void set(ISpan * _span);
+    void setown(ISpan * _span);
+
+private:
+    Owned<ISpan> span;
+    ISpan * prevSpan = nullptr;
 };
 
 extern jlib_decl IProperties * getClientHeaders(const ISpan * span);
@@ -180,6 +188,7 @@ interface ITraceManager : extends IInterface
     virtual bool isTracingEnabled() const = 0;
  };
 
+extern jlib_decl ISpan * queryNullSpan();
 extern jlib_decl ISpan * getNullSpan();
 extern jlib_decl void initTraceManager(const char * componentName, const IPropertyTree * componentConfig, const IPropertyTree * globalConfig);
 extern jlib_decl ITraceManager & queryTraceManager();