Skip to content

Commit ecb8f4a

Browse files
author
Ignacio Bonafonte
authored
Implement a Stack in the TaskLocal for automatic context handling (#284)
Reactivate TaskLocal context, but use OS_ACTIVITY_CURRENT value logic to know if we are inside a Task or not ( inside a task, createActivityContext() always returns nil, and OS_ACTIVITY_CURRENT always returns 0 Added some tests to test current working and failing scenarios
1 parent cadbcbf commit ecb8f4a

File tree

7 files changed

+465
-41
lines changed

7 files changed

+465
-41
lines changed

‎Sources/OpenTelemetryApi/Context/ActivityContextManager.swift‎

Lines changed: 82 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,13 @@ private let OS_ACTIVITY_CURRENT = unsafeBitCast(dlsym(UnsafeMutableRawPointer(bi
1717
class ActivityContextManager: ContextManager {
1818
static let instance = ActivityContextManager()
1919
#if canImport(_Concurrency)
20+
#if swift(<5.5.2)
2021
@available(macOS 12.0, iOS 15.0, tvOS 15.0, *)
2122
static let taskLocalContextManager = TaskLocalContextManager.instance
23+
#else
24+
@available(macOS 10.15, iOS 13.0, tvOS 13.0, *)
25+
static let taskLocalContextManager = TaskLocalContextManager.instance
26+
#endif
2227
#endif
2328

2429
let rlock = NSRecursiveLock()
@@ -38,45 +43,80 @@ class ActivityContextManager: ContextManager {
3843
var contextMap = [os_activity_id_t: [String: AnyObject]]()
3944

4045
func getCurrentContextValue(forKey key: OpenTelemetryContextKeys) -> AnyObject? {
41-
#if canImport(_Concurrency)
42-
if #available(macOS 12.0, iOS 15.0, tvOS 15.0, *) {
43-
// If running with task local, use first its stored value
44-
if let contextValue = ActivityContextManager.taskLocalContextManager.getCurrentContextValue(forKey: key) {
45-
return contextValue
46-
}
47-
}
48-
#endif
4946
var parentIdent: os_activity_id_t = 0
5047
let activityIdent = os_activity_get_identifier(OS_ACTIVITY_CURRENT, &parentIdent)
5148
var contextValue: AnyObject?
52-
rlock.lock()
53-
guard let context = contextMap[activityIdent] ?? contextMap[parentIdent] else {
49+
if activityIdent != 0 {
50+
rlock.lock()
51+
guard let context = contextMap[activityIdent] ?? contextMap[parentIdent] else {
52+
rlock.unlock()
53+
return nil
54+
}
55+
contextValue = context[key.rawValue]
5456
rlock.unlock()
55-
return nil
57+
return contextValue
58+
} else {
59+
// If activityIdent == 0, it means no active Span or we are inside an Task
60+
#if canImport(_Concurrency)
61+
#if swift(<5.5.2)
62+
if #available(macOS 12.0, iOS 15.0, tvOS 15.0, *) {
63+
if let contextValue = ActivityContextManager.taskLocalContextManager.getCurrentContextValue(forKey: key) {
64+
return contextValue
65+
}
66+
}
67+
#else
68+
if #available(macOS 10.15, iOS 13.0, tvOS 13.0, *) {
69+
if let contextValue = ActivityContextManager.taskLocalContextManager.getCurrentContextValue(forKey: key) {
70+
return contextValue
71+
}
72+
}
73+
#endif
74+
#endif
5675
}
57-
contextValue = context[key.rawValue]
58-
rlock.unlock()
59-
return contextValue
76+
return nil
6077
}
6178

6279
func setCurrentContextValue(forKey key: OpenTelemetryContextKeys, value: AnyObject) {
6380
var parentIdent: os_activity_id_t = 0
6481
var activityIdent = os_activity_get_identifier(OS_ACTIVITY_CURRENT, &parentIdent)
65-
rlock.lock()
66-
if contextMap[activityIdent] == nil || contextMap[activityIdent]?[key.rawValue] != nil {
82+
if activityIdent != 0 {
83+
// We are inside an activity, it can be an activity created by us for a span context or another independent activty
84+
// We are surely not inside a Task
85+
rlock.lock()
86+
if contextMap[activityIdent] == nil || contextMap[activityIdent]?[key.rawValue] != nil {
87+
var scope: os_activity_scope_state_s
88+
(activityIdent, scope) = createActivityContext()
89+
contextMap[activityIdent] = [String: AnyObject]()
90+
objectScope.setObject(ScopeElement(scope: scope), forKey: value)
91+
}
92+
contextMap[activityIdent]?[key.rawValue] = value
93+
rlock.unlock()
94+
} else {
6795
var scope: os_activity_scope_state_s
6896
(activityIdent, scope) = createActivityContext()
69-
contextMap[activityIdent] = [String: AnyObject]()
70-
objectScope.setObject(ScopeElement(scope: scope), forKey: value)
71-
}
72-
contextMap[activityIdent]?[key.rawValue] = value
97+
if activityIdent == 0 {
98+
// If activityIdent == 0, means we are inside a Task, because we cannot create an activity, set the context inside the task
7399
#if canImport(_Concurrency)
74-
if #available(macOS 12.0, iOS 15.0, tvOS 15.0, *) {
75-
// If running with task local, set the value after the activity, so activity is not empty
76-
ActivityContextManager.taskLocalContextManager.setCurrentContextValue(forKey: key, value: value)
77-
}
100+
#if swift(<5.5.2)
101+
if #available(macOS 12.0, iOS 15.0, tvOS 15.0, *) {
102+
ActivityContextManager.taskLocalContextManager.setCurrentContextValue(forKey: key, value: value)
103+
}
104+
#else
105+
if #available(macOS 10.15, iOS 13.0, tvOS 13.0, *) {
106+
ActivityContextManager.taskLocalContextManager.setCurrentContextValue(forKey: key, value: value)
107+
}
108+
109+
#endif
78110
#endif
79-
rlock.unlock()
111+
} else {
112+
// We could create the activity so we store the context in the activity map
113+
rlock.lock()
114+
contextMap[activityIdent] = [String: AnyObject]()
115+
objectScope.setObject(ScopeElement(scope: scope), forKey: value)
116+
contextMap[activityIdent]?[key.rawValue] = value
117+
rlock.unlock()
118+
}
119+
}
80120
}
81121

82122
func createActivityContext() -> (os_activity_id_t, os_activity_scope_state_s) {
@@ -93,15 +133,26 @@ class ActivityContextManager: ContextManager {
93133
var scope = scope.scope
94134
os_activity_scope_leave(&scope)
95135
objectScope.removeObject(forKey: value)
96-
}
136+
} else {
97137
#if canImport(_Concurrency)
98-
if #available(macOS 12.0, iOS 15.0, tvOS 15.0, *) {
138+
#if swift(<5.5.2)
139+
if #available(macOS 12.0, iOS 15.0, tvOS 15.0, *) {
140+
// If there is a parent activity, set its content as the task local
141+
ActivityContextManager.taskLocalContextManager.removeContextValue(forKey: key, value: value)
142+
if let currentContext = self.getCurrentContextValue(forKey: key) {
143+
ActivityContextManager.taskLocalContextManager.setCurrentContextValue(forKey: key, value: currentContext)
144+
}
145+
}
146+
#else
99147
// If there is a parent activity, set its content as the task local
100-
ActivityContextManager.taskLocalContextManager.removeContextValue(forKey: key, value: value)
101-
if let currentContext = self.getCurrentContextValue(forKey: key) {
102-
ActivityContextManager.taskLocalContextManager.setCurrentContextValue(forKey: key, value: currentContext)
148+
if #available(macOS 10.15, iOS 13.0, tvOS 13.0, *) {
149+
ActivityContextManager.taskLocalContextManager.removeContextValue(forKey: key, value: value)
150+
if let currentContext = self.getCurrentContextValue(forKey: key) {
151+
ActivityContextManager.taskLocalContextManager.setCurrentContextValue(forKey: key, value: currentContext)
152+
}
103153
}
104-
}
105154
#endif
155+
#endif
156+
}
106157
}
107158
}

‎Sources/OpenTelemetryApi/Context/TaskLocalContextManager.swift‎

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,105 @@
66
import Foundation
77

88
#if canImport(_Concurrency)
9+
#if swift(>=5.5.2)
10+
@available(macOS 10.15, iOS 13.0, tvOS 13.0, *)
11+
enum ContextManagement {
12+
static let spanRLock = NSRecursiveLock()
13+
static let baggageRLock = NSRecursiveLock()
14+
15+
@TaskLocal
16+
private static var spans = [Span]()
17+
18+
public static func getSpan() -> Span? {
19+
spanRLock.lock()
20+
defer { spanRLock.unlock() }
21+
return spans.last
22+
}
23+
24+
public static func setSpan(span: Span) {
25+
spanRLock.lock()
26+
defer { spanRLock.unlock() }
27+
var aux = spans
28+
aux.append(span)
29+
_spans = TaskLocal(wrappedValue: aux)
30+
}
31+
32+
public static func removeSpan() {
33+
spanRLock.lock()
34+
defer { spanRLock.unlock() }
35+
guard !spans.isEmpty else {
36+
return
37+
}
38+
var aux = spans
39+
aux.removeLast()
40+
_spans = TaskLocal(wrappedValue: aux)
41+
}
42+
43+
@TaskLocal
44+
private static var baggages = [Baggage]()
45+
46+
public static func getBaggage() -> Baggage? {
47+
baggageRLock.lock()
48+
defer { baggageRLock.unlock() }
49+
return baggages.last
50+
}
51+
52+
public static func setBaggage(baggage: Baggage) {
53+
baggageRLock.lock()
54+
defer { baggageRLock.unlock() }
55+
var aux = baggages
56+
aux.append(baggage)
57+
_baggages = TaskLocal(wrappedValue: aux)
58+
}
59+
60+
public static func removeBaggage() {
61+
baggageRLock.lock()
62+
defer { baggageRLock.unlock() }
63+
guard !baggages.isEmpty else {
64+
return
65+
}
66+
var aux = baggages
67+
aux.removeLast()
68+
_baggages = TaskLocal(wrappedValue: aux)
69+
}
70+
}
71+
72+
@available(macOS 10.15, iOS 13.0, tvOS 13.0, *)
73+
class TaskLocalContextManager: ContextManager {
74+
static let instance = TaskLocalContextManager()
75+
76+
func getCurrentContextValue(forKey key: OpenTelemetryContextKeys) -> AnyObject? {
77+
switch key {
78+
case .span:
79+
return ContextManagement.getSpan()
80+
case .baggage:
81+
return ContextManagement.getBaggage()
82+
}
83+
}
84+
85+
func setCurrentContextValue(forKey key: OpenTelemetryContextKeys, value: AnyObject) {
86+
switch key {
87+
case .span:
88+
if let span = value as? Span {
89+
ContextManagement.setSpan(span: span)
90+
}
91+
case .baggage:
92+
if let baggage = value as? Baggage {
93+
ContextManagement.setBaggage(baggage: baggage)
94+
}
95+
}
96+
}
97+
98+
func removeContextValue(forKey key: OpenTelemetryContextKeys, value: AnyObject) {
99+
switch key {
100+
case .span:
101+
ContextManagement.removeSpan()
102+
case .baggage:
103+
ContextManagement.removeBaggage()
104+
}
105+
}
106+
}
107+
#else
9108
@available(macOS 12.0, iOS 15.0, tvOS 15.0, *)
10109
enum ContextManagement {
11110
@TaskLocal
@@ -57,3 +156,4 @@ class TaskLocalContextManager: ContextManager {
57156
}
58157
}
59158
#endif
159+
#endif

‎Sources/OpenTelemetryApi/Trace/PropagatedSpanBuilder.swift‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,19 +10,24 @@ class PropagatedSpanBuilder: SpanBuilder {
1010
private var tracer: Tracer
1111
private var isRootSpan: Bool = false
1212
private var spanContext: SpanContext?
13+
private var spanName: String
1314

14-
init(tracer: Tracer, spanName: String) {
15+
16+
init(tracer: Tracer, spanName: String = "") {
1517
self.tracer = tracer
18+
self.spanName = spanName
1619
}
1720

1821
@discardableResult public func startSpan() -> Span {
1922
if spanContext == nil, !isRootSpan {
2023
spanContext = OpenTelemetry.instance.contextProvider.activeSpan?.context
2124
}
22-
return PropagatedSpan(context: spanContext ?? SpanContext.create(traceId: TraceId.random(),
25+
let span = PropagatedSpan(context: spanContext ?? SpanContext.create(traceId: TraceId.random(),
2326
spanId: SpanId.random(),
2427
traceFlags: TraceFlags(),
2528
traceState: TraceState()))
29+
span.name = spanName
30+
return span
2631
}
2732

2833
@discardableResult public func setParent(_ parent: Span) -> Self {

‎Tests/OpenTelemetryApiTests/Baggage/DefaultBaggageManagerTests.swift‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -28,11 +28,9 @@ class DefaultBaggageManagerTests: XCTestCase {
2828
let baggage = TestBaggage()
2929

3030
override func tearDown() {
31-
if defaultBaggageManager.getCurrentBaggage() != nil {
32-
XCTAssert(false, "Test must clean baggage context")
33-
}
31+
XCTAssertNil(defaultBaggageManager.getCurrentBaggage(), "Test must clean baggage context")
3432
}
35-
33+
3634
func testBuilderMethod() {
3735
let builder = defaultBaggageManager.baggageBuilder()
3836
XCTAssertEqual(builder.build().getEntries().count, 0)
@@ -62,6 +60,7 @@ class DefaultBaggageManagerTests: XCTestCase {
6260
let semaphore = DispatchSemaphore(value: 0)
6361
let semaphore2 = DispatchSemaphore(value: 0)
6462
DispatchQueue.global().async {
63+
XCTAssert(self.defaultBaggageManager.getCurrentBaggage() === self.baggage)
6564
semaphore2.signal()
6665
semaphore.wait()
6766
XCTAssert(self.defaultBaggageManager.getCurrentBaggage() === self.baggage)

0 commit comments

Comments
 (0)