From 27e3fecca5ccdd6cac543875e3f5a51cb4245b28 Mon Sep 17 00:00:00 2001
From: Ignacio Bonafonte <nacho.bonafontearruga@datadoghq.com>
Date: Thu, 3 Aug 2023 19:08:19 +0200
Subject: [PATCH] Fix active context is removed wrongly (#448)

The current context must be removed only if it is the active span
Add a test for the non active span end
---
 .../Context/ActivityContextManager.swift         |  8 ++++++--
 .../Context/ActivityContextManagerTests.swift    | 16 ++++++++++++++++
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/Sources/OpenTelemetryApi/Context/ActivityContextManager.swift b/Sources/OpenTelemetryApi/Context/ActivityContextManager.swift
index 1b7d40fa..2f5596d5 100644
--- 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
index d579a45e..28f40ad5 100644
--- 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() {