Merge branch 'main' into otlp-version-fix

open-telemetry · Aug 7, 2023 · bb17394 · bb17394
2 parents abc4ad1 + 27e3fec
commit bb17394
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 2 deletions.
diff --git a/Sources/OpenTelemetryApi/Context/ActivityContextManager.swift b/Sources/OpenTelemetryApi/Context/ActivityContextManager.swift
@@ -70,9 +70,13 @@ class ActivityContextManager: ContextManager {
     func removeContextValue(forKey key: OpenTelemetryContextKeys, value: AnyObject) {
         let activityIdent = os_activity_get_identifier(OS_ACTIVITY_CURRENT, nil)
         rlock.lock()
-        contextMap[activityIdent]?[key.rawValue] = nil
-        if contextMap[activityIdent]?.isEmpty ?? false {
+        if let currentValue = contextMap[activityIdent]?[key.rawValue],
+           currentValue === value
+        {
+         contextMap[activityIdent]?[key.rawValue] = nil
+          if contextMap[activityIdent]?.isEmpty ?? false {
             contextMap[activityIdent] = nil
+          }
         }
         if let scope = objectScope.object(forKey: value) {
             var scope = scope.scope

diff --git a/Tests/OpenTelemetryApiTests/Context/ActivityContextManagerTests.swift b/Tests/OpenTelemetryApiTests/Context/ActivityContextManagerTests.swift
@@ -77,6 +77,22 @@ class ActivityContextManagerTests: XCTestCase {
         XCTAssert(activeSpan === parentSpan)
     }
 
+    func testContextPropagationTwoSequentialChildSpans() {
+        let parentSpan = defaultTracer.spanBuilder(spanName: "Parent").startSpan()
+        OpenTelemetry.instance.contextProvider.setActiveSpan(parentSpan)
+
+        let child1 = defaultTracer.spanBuilder(spanName: "child1").startSpan()
+        child1.end()
+
+        let child2 = defaultTracer.spanBuilder(spanName: "child2").startSpan()
+        child2.end()
+
+        parentSpan.end()
+
+        XCTAssertEqual(parentSpan.context.traceId, child1.context.traceId)
+        XCTAssertEqual(parentSpan.context.traceId, child2.context.traceId)
+    }
+
     #if swift(>=5.5.2)
     @available(macOS 10.15, iOS 13.0, tvOS 13.0, *)
     func testStartAndEndSpanInAsyncTask() {