Fix span nesting on iOS

example/src/WixApp.tsx

@@ -1,4 +1,4 @@
-import React from 'react';
+import React, { useState } from 'react';
 import { View, Text, Button } from 'react-native';
 import MainScreen from './screens/MainScreen';
 import ErrorScreen from './screens/ErrorScreen';
@@ -11,6 +11,8 @@ import {
 } from '@datadog/mobile-react-native-navigation';
 
 import styles from './screens/styles';
+import { DdTrace } from '@datadog/mobile-react-native';
+import TraceScreen from './screens/TraceScreen';
 
 const viewPredicate: ViewNamePredicate = (
     _event: ComponentDidAppearEvent,
@@ -37,6 +39,7 @@ function registerScreens() {
     Navigation.registerComponent('Home', () => HomeScreen);
     Navigation.registerComponent('Main', () => MainScreen);
     Navigation.registerComponent('Error', () => ErrorScreen);
+    Navigation.registerComponent('Trace', () => TraceScreen);
     Navigation.registerComponent('About', () => AboutScreen);
 }
 
@@ -64,6 +67,15 @@ const HomeScreen = props => {
                 }}
             />
             <View style={{ marginTop: 20 }} />
+            <Button
+                title="Trace"
+                onPress={() => {
+                    Navigation.push(props.componentId, {
+                        component: { name: 'Trace' }
+                    });
+                }}
+            />
+            <View style={{ marginTop: 20 }} />
             <Button
                 title="About"
                 onPress={() => {

example/src/screens/TraceScreen.tsx

@@ -0,0 +1,37 @@
+import { useState } from 'react';
+import { View, Text, Button } from 'react-native';
+
+import styles from './styles';
+import { DdTrace } from '@datadog/mobile-react-native';
+
+const TraceScreen = props => {
+    const [spanIds, setSpanIds] = useState<string[]>([]);
+
+    return (
+        <View style={styles.defaultScreen}>
+            <Text style={{ marginBottom: 20 }}>
+                Trace test
+            </Text>
+            <View style={{ marginTop: 20 }} />
+            <Button
+                title="Start Span"
+                onPress={async () => {
+                    const newSpanId = await DdTrace.startSpan("span_operation_" +   Math.floor(100 + Math.random() * 900));
+                    spanIds.push(newSpanId);
+                    setSpanIds(spanIds);
+                }}
+            />
+            <View style={{ marginTop: 20 }} />
+            <Button
+                title="Finish Span"
+                onPress={() => {
+                    const lastSpan = spanIds.pop();
+                    lastSpan && DdTrace.finishSpan(lastSpan);
+                    setSpanIds(spanIds)
+                }}
+            />
+        </View>
+    );
+};
+
+export default TraceScreen;

packages/core/ios/Sources/DdTraceImplementation.swift

@@ -4,14 +4,20 @@
  * Copyright 2016-Present Datadog, Inc.
  */
 
-import Foundation
 import DatadogTrace
+import Foundation
 
 @objc
 public class DdTraceImplementation: NSObject {
     private lazy var tracer: OTTracer = tracerProvider()
     private let tracerProvider: () -> OTTracer
-    private(set) var spanDictionary = [NSString: OTSpan]()
+
+    private(set) var spansById: [String: OTSpan] = [:]
+    private(set) var spanStack: [String] = []
+
+    private var activeSpan: OTSpan? {
+        spanStack.last.flatMap { spansById[$0] }
+    }
 
     internal init(_ tracerProvider: @escaping () -> OTTracer) {
         self.tracerProvider = tracerProvider
@@ -23,35 +29,56 @@ public class DdTraceImplementation: NSObject {
     }
 
     @objc
-    public func startSpan(operation: String, context: NSDictionary, timestampMs: Double, resolve:RCTPromiseResolveBlock, reject:RCTPromiseRejectBlock) -> Void {
-        let id = UUID().uuidString as NSString
+    public func startSpan(
+        operation: String, context: NSDictionary, timestampMs: Double,
+        resolve: RCTPromiseResolveBlock, reject: RCTPromiseRejectBlock
+    ) {
+        objc_sync_enter(self)
+        defer { objc_sync_exit(self) }
+
+        let id = UUID().uuidString
         let timeIntervalSince1970: TimeInterval = timestampMs / 1_000
         let startDate = Date(timeIntervalSince1970: timeIntervalSince1970)
 
-        objc_sync_enter(self)
         let span = tracer.startSpan(
             operationName: operation,
-            childOf: nil,
+            childOf: activeSpan?.context,
             tags: castAttributesToSwift(context).mergeWithGlobalAttributes(),
             startTime: startDate
         )
+
         span.setActive()
-        spanDictionary[id] = span
-        objc_sync_exit(self)
+        spansById[id] = span
+        spanStack.append(id)
 
         resolve(id)
     }
 
     @objc
-    public func finishSpan(spanId: NSString, context: NSDictionary, timestampMs: Double, resolve:RCTPromiseResolveBlock, reject:RCTPromiseRejectBlock) -> Void {
+    public func finishSpan(
+        spanId: NSString, context: NSDictionary, timestampMs: Double,
+        resolve: RCTPromiseResolveBlock, reject: RCTPromiseRejectBlock
+    ) {
         objc_sync_enter(self)
-        let optionalSpan = spanDictionary.removeValue(forKey: spanId)
-        if let span = optionalSpan {
-            set(tags: castAttributesToSwift(context).mergeWithGlobalAttributes(), to: span)
-            let timeIntervalSince1970: TimeInterval = timestampMs / 1_000
-            span.finish(at: Date(timeIntervalSince1970: timeIntervalSince1970))
+        defer { objc_sync_exit(self) }
+
+        guard let span = spansById.removeValue(forKey: spanId as String) else {
+            resolve(nil)
+            return
+        }
+
+        set(tags: castAttributesToSwift(context).mergeWithGlobalAttributes(), to: span)
+        let timeIntervalSince1970: TimeInterval = timestampMs / 1_000
+        span.finish(at: Date(timeIntervalSince1970: timeIntervalSince1970))
+
+        if let idx = spanStack.lastIndex(of: spanId as String) {
+            let wasTop = (idx == spanStack.count - 1)
+            spanStack.remove(at: idx)
+
+            if wasTop, let prev = activeSpan {
+                prev.setActive()
+            }
         }
-        objc_sync_exit(self)
 
         resolve(nil)
     }

packages/core/ios/Tests/DdTraceTests.swift

@@ -90,20 +90,22 @@ internal class DdTraceTests: XCTestCase {
             resolve: mockResolve,
             reject: mockReject
         )
-        let spanID = lastResolveValue as! NSString
+        let spanID = lastResolveValue as! String
 
         XCTAssertNotNil(spanID)
-        XCTAssertEqual(Array(tracer.spanDictionary.keys), [spanID])
+        XCTAssertEqual(Array(tracer.spansById.keys), [spanID])
+        XCTAssertEqual(Array(tracer.spanStack), [spanID])
         let startedSpan = try XCTUnwrap(mockNativeTracer.startedSpans.last)
         XCTAssertEqual(startedSpan.finishTime, MockSpan.unfinished)
 
         let spanDuration: TimeInterval = 10.042
         let spanDurationMs = spanDuration * 1_000
         let finishTimestampMs = timestampMs + spanDurationMs
         let finishingContext = NSDictionary(dictionary: ["last_key": "last_value"])
-        tracer.finishSpan(spanId: spanID, context: finishingContext, timestampMs: finishTimestampMs, resolve: mockResolve, reject: mockReject)
+        tracer.finishSpan(spanId: spanID as NSString, context: finishingContext, timestampMs: finishTimestampMs, resolve: mockResolve, reject: mockReject)
 
-        XCTAssertEqual(Array(tracer.spanDictionary.keys), [])
+        XCTAssertEqual(Array(tracer.spansById.keys), [])
+        XCTAssertEqual(Array(tracer.spanStack), [])
         XCTAssertEqual(
             startedSpan.finishTime!.timeIntervalSince1970, // swiftlint:disable:this force_unwrapping
             (startDate + spanDuration).timeIntervalSince1970,
@@ -124,7 +126,8 @@ internal class DdTraceTests: XCTestCase {
             reject: mockReject
         )
 
-        XCTAssertEqual(tracer.spanDictionary.count, 1)
+        XCTAssertEqual(tracer.spansById.count, 1)
+        XCTAssertEqual(tracer.spanStack.count, 1)
 
         XCTAssertNoThrow(
             tracer.finishSpan(
@@ -136,7 +139,8 @@ internal class DdTraceTests: XCTestCase {
             )
         )
 
-        XCTAssertEqual(tracer.spanDictionary.count, 1)
+        XCTAssertEqual(tracer.spansById.count, 1)
+        XCTAssertEqual(tracer.spanStack.count, 1)
     }
 
     func testTracingConcurrently() {
@@ -158,7 +162,8 @@ internal class DdTraceTests: XCTestCase {
         }
 
         XCTAssertEqual(mockNativeTracer.startedSpans.count, iterationCount, "\(mockNativeTracer.startedSpans)")
-        XCTAssertEqual(tracer.spanDictionary.count, 0)
+        XCTAssertEqual(tracer.spansById.count, 0)
+        XCTAssertEqual(tracer.spanStack.count, 0)
     }
 }
 
@@ -205,8 +210,15 @@ private class MockSpan: OTSpan {
         self.finishTime = time
     }
 
+    private final class Ctx: OTSpanContext {
+        func forEachBaggageItem(callback: (String, String) -> Bool) {
+            // No baggage in tests → do nothing
+        }
+    }
+
     // swiftlint:disable unavailable_function
-    var context: OTSpanContext { fatalError("Should not be called") }
+    var context: OTSpanContext = Ctx()
+
     func tracer() -> OTTracer {
         fatalError("Should not be called")
     }