Updates URLSessionInstrumentation to check task.state before setting …

…FakeDelegaet (#539)
open-telemetry · May 7, 2024 · 60eac93 · 60eac93
1 parent 62772bd
commit 60eac93
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 2 deletions.
diff --git a/Sources/Instrumentation/URLSession/URLSessionInstrumentation.swift b/Sources/Instrumentation/URLSession/URLSessionInstrumentation.swift
@@ -600,14 +600,15 @@ public class URLSessionInstrumentation {
 
             if #available(OSX 13.0, iOS 16.0, watchOS 9.0, tvOS 16.0, *) {
                 guard Task.basePriority != nil else {
+                    // If not inside a Task basePriority is nil
                     return
                 }
                 let instrumentedRequest = URLSessionLogger.processAndLogRequest(request, sessionTaskId: taskId, instrumentation: self, shouldInjectHeaders: true)
                 task.setValue(instrumentedRequest, forKey: "currentRequest")
                 self.setIdKey(value: taskId, for: task)
 
-                // If not inside a Task basePriority is nil
-                if task.delegate == nil {
+
+                if task.delegate == nil && task.state != .running {
                     task.delegate = FakeDelegate()
                 }
             }

diff --git a/Tests/InstrumentationTests/URLSessionTests/URLSessionInstrumentationTests.swift b/Tests/InstrumentationTests/URLSessionTests/URLSessionInstrumentationTests.swift
@@ -304,6 +304,22 @@ class URLSessionInstrumentationTests: XCTestCase {
         XCTAssertNotNil(URLSessionInstrumentationTests.requestCopy?.allHTTPHeaderFields?[W3CTraceContextPropagator.traceparent])
     }
 
+    public func testDataTaskWithUrlBlock_doesNotCrashWhenResumed_whenDataTaskIsRunning() {
+        let url = URL(string: "http://localhost:33333/success")!
+        let dataTask = URLSession.shared.dataTask(with: url) { _, _, _ in
+            URLSessionInstrumentationTests.semaphore.signal()
+        }
+
+        dataTask.resume()
+        Task.detached {
+            dataTask.resume()
+        }
+
+        URLSessionInstrumentationTests.semaphore.wait()
+        XCTAssertTrue(URLSessionInstrumentationTests.checker.createdRequestCalled)
+        XCTAssertNotNil(URLSessionInstrumentationTests.requestCopy?.allHTTPHeaderFields?[W3CTraceContextPropagator.traceparent])
+    }
+
     public func testDownloadTaskWithUrlBlock() {
         let url = URL(string: "http://localhost:33333/success")!