Merge pull request #230 from nachoBonafonte/TaskLocal-Context

Add support for automatic context handling when using Swift structured concurrency.
Create a TaskLocal context manager that can handle async/await tasks and also use it from the ActivityContext manager, so we can mix both behaviours.
Refactor a bit the classes and protocols for consistency.

Author: Ignacio Bonafonte

Sources/OpenTelemetryApi/Context/ActivityContextManager.swift

@@ -16,6 +16,10 @@ private let OS_ACTIVITY_CURRENT = unsafeBitCast(dlsym(UnsafeMutableRawPointer(bi
 
 class ActivityContextManager: ContextManager {
     static let instance = ActivityContextManager()
+    #if swift(>=5.5)
+    @available(macOS 12.0, iOS 15.0, tvOS 15.0, *)
+    static let taskLocalContextManager = TaskLocalContextManager.instance
+    #endif
 
     let rlock = NSRecursiveLock()
 
@@ -24,8 +28,7 @@ class ActivityContextManager: ContextManager {
             self.scope = scope
         }
 
-        deinit {
-        }
+        deinit {}
 
         var scope: os_activity_scope_state_s
     }
@@ -34,7 +37,15 @@ class ActivityContextManager: ContextManager {
 
     var contextMap = [os_activity_id_t: [String: AnyObject]]()
 
-    func getCurrentContextValue(forKey key: String) -> AnyObject? {
+    func getCurrentContextValue(forKey key: OpenTelemetryContextKeys) -> AnyObject? {
+        #if swift(>=5.5)
+        if #available(macOS 12.0, iOS 15.0, tvOS 15.0, *) {
+            // If running with task local, use first its stored value
+            if let contextValue = ActivityContextManager.taskLocalContextManager.getCurrentContextValue(forKey: key) {
+                return contextValue
+            }
+        }
+        #endif
         var parentIdent: os_activity_id_t = 0
         let activityIdent = os_activity_get_identifier(OS_ACTIVITY_CURRENT, &parentIdent)
         var contextValue: AnyObject?
@@ -43,22 +54,28 @@ class ActivityContextManager: ContextManager {
             rlock.unlock()
             return nil
         }
-        contextValue = context[key]
+        contextValue = context[key.rawValue]
         rlock.unlock()
         return contextValue
     }
 
-    func setCurrentContextValue(forKey key: String, value: AnyObject) {
+    func setCurrentContextValue(forKey key: OpenTelemetryContextKeys, value: AnyObject) {
         var parentIdent: os_activity_id_t = 0
         var activityIdent = os_activity_get_identifier(OS_ACTIVITY_CURRENT, &parentIdent)
         rlock.lock()
-        if contextMap[activityIdent] == nil || contextMap[activityIdent]?[key] != nil {
+        if contextMap[activityIdent] == nil || contextMap[activityIdent]?[key.rawValue] != nil {
             var scope: os_activity_scope_state_s
             (activityIdent, scope) = createActivityContext()
             contextMap[activityIdent] = [String: AnyObject]()
             objectScope.setObject(ScopeElement(scope: scope), forKey: value)
         }
-        contextMap[activityIdent]?[key] = value
+        contextMap[activityIdent]?[key.rawValue] = value
+        #if swift(>=5.5)
+        if #available(macOS 12.0, iOS 15.0, tvOS 15.0, *) {
+            // If running with task local, set the value after the activity, so activity is not empty
+            ActivityContextManager.taskLocalContextManager.setCurrentContextValue(forKey: key, value: value)
+        }
+        #endif
         rlock.unlock()
     }
 
@@ -71,11 +88,20 @@ class ActivityContextManager: ContextManager {
         return (currentActivityId, activityState)
     }
 
-    func removeContextValue(forKey key: String, value: AnyObject) {
+    func removeContextValue(forKey key: OpenTelemetryContextKeys, value: AnyObject) {
         if let scope = objectScope.object(forKey: value) {
             var scope = scope.scope
             os_activity_scope_leave(&scope)
             objectScope.removeObject(forKey: value)
         }
+        #if swift(>=5.5)
+        if #available(macOS 12.0, iOS 15.0, tvOS 15.0, *) {
+            // If there is a parent activity, set its content as the task local
+            ActivityContextManager.taskLocalContextManager.removeContextValue(forKey: key, value: value)
+            if let currentContext = self.getCurrentContextValue(forKey: key) {
+                ActivityContextManager.taskLocalContextManager.setCurrentContextValue(forKey: key, value: currentContext)
+            }
+        }
+        #endif
     }
 }

Sources/OpenTelemetryApi/Context/ContextManager.swift

@@ -6,7 +6,7 @@
 import Foundation
 
 public protocol ContextManager: AnyObject {
-    func getCurrentContextValue(forKey: String) -> AnyObject?
-    func setCurrentContextValue(forKey: String, value: AnyObject)
-    func removeContextValue(forKey: String, value: AnyObject)
+    func getCurrentContextValue(forKey: OpenTelemetryContextKeys) -> AnyObject?
+    func setCurrentContextValue(forKey: OpenTelemetryContextKeys, value: AnyObject)
+    func removeContextValue(forKey: OpenTelemetryContextKeys, value: AnyObject)
 }

Sources/OpenTelemetryApi/Context/OpenTelemetryContextProvider.swift

@@ -7,20 +7,17 @@ import Foundation
 import os.activity
 
 /// Keys used by Opentelemetry to store values in the Context
-internal struct OpenTelemetryContextKeys {
-    static let span = "opentelemetrycontext.span"
-    static let baggage    = "opentelemetrycontext.baggage"
+public enum OpenTelemetryContextKeys: String {
+    case span
+    case baggage
 }
 
-
 public struct OpenTelemetryContextProvider {
-
     var contextManager: ContextManager
 
-
     /// Returns the Span from the current context
     public var activeSpan: Span? {
-        return contextManager.getCurrentContextValue(forKey: OpenTelemetryContextKeys.span) as? Span
+        return contextManager.getCurrentContextValue(forKey: .span) as? Span
     }
 
     /// Returns the Baggage from the current context
@@ -30,7 +27,7 @@ public struct OpenTelemetryContextProvider {
 
     /// Sets the span as the activeSpan for the current context
     /// - Parameter span: the Span to be set to the current context
-    public func setActiveSpan(_ span: Span)  {
+    public func setActiveSpan(_ span: Span) {
         contextManager.setCurrentContextValue(forKey: OpenTelemetryContextKeys.span, value: span)
     }
 

Sources/OpenTelemetryApi/Context/TaskLocalContextManager.swift

@@ -0,0 +1,59 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import Foundation
+
+#if swift(>=5.5)
+@available(macOS 12.0, iOS 15.0, tvOS 15.0, *)
+enum ContextManagement {
+    @TaskLocal
+    static var span: Span?
+    public static func setSpan(wrappedSpan: TaskLocal<Span?>) {
+        _span = wrappedSpan
+    }
+
+    @TaskLocal
+    static var baggage: Baggage?
+    public static func setBaggage(wrappedBaggage: TaskLocal<Baggage?>) {
+        _baggage = wrappedBaggage
+    }
+}
+
+@available(macOS 12.0, iOS 15.0, tvOS 15.0, *)
+class TaskLocalContextManager: ContextManager {
+    static let instance = TaskLocalContextManager()
+
+    func getCurrentContextValue(forKey key: OpenTelemetryContextKeys) -> AnyObject? {
+        switch key {
+            case .span:
+                return ContextManagement.span
+            case .baggage:
+                return ContextManagement.baggage
+        }
+    }
+
+    func setCurrentContextValue(forKey key: OpenTelemetryContextKeys, value: AnyObject) {
+        switch key {
+            case .span:
+                if let span = value as? Span {
+                    ContextManagement.setSpan(wrappedSpan: TaskLocal(wrappedValue: span))
+                }
+            case .baggage:
+                if let baggage = value as? Baggage {
+                    ContextManagement.setBaggage(wrappedBaggage: TaskLocal(wrappedValue: baggage))
+                }
+        }
+    }
+
+    func removeContextValue(forKey key: OpenTelemetryContextKeys, value: AnyObject) {
+        switch key {
+            case .span:
+                ContextManagement.setSpan(wrappedSpan: TaskLocal(wrappedValue: nil))
+            case .baggage:
+                ContextManagement.setBaggage(wrappedBaggage: TaskLocal(wrappedValue: nil))
+        }
+    }
+}
+#endif