async await http instrumentation (#849)

Author: martin308

Sources/Instrumentation/URLSession/URLSessionInstrumentation.swift

@@ -721,6 +721,45 @@ public class URLSessionInstrumentation {
         requestMap[taskId]?.setRequest(request)
       }
 
+      // For iOS 15+/macOS 12+, handle async/await methods differently
+      if #available(macOS 12.0, iOS 15.0, tvOS 15.0, watchOS 8.0, *) {
+        // Check if we can determine if this is an async/await call
+        // For iOS 15/macOS 12, we can't use Task.basePriority, so we check other indicators
+        var isAsyncContext = false
+        
+        if #available(OSX 13.0, iOS 16.0, watchOS 9.0, tvOS 16.0, *) {
+          isAsyncContext = Task.basePriority != nil
+        } else {
+          // For iOS 15/macOS 12, check if the task has no delegate and no session delegate
+          // This is a heuristic that works for async/await methods
+          isAsyncContext = task.delegate == nil && 
+                          (task.value(forKey: "session") as? URLSession)?.delegate == nil &&
+                          task.state != .running
+        }
+        
+        if isAsyncContext {
+          // This is likely an async/await call
+          let instrumentedRequest = URLSessionLogger.processAndLogRequest(request,
+                                                                        sessionTaskId: taskId,
+                                                                        instrumentation: self,
+                                                                        shouldInjectHeaders: true)
+          if let instrumentedRequest {
+            task.setValue(instrumentedRequest, forKey: "currentRequest")
+          }
+          self.setIdKey(value: taskId, for: task)
+          
+          // For async/await methods, we need to ensure the delegate is set
+          // to capture the completion, but only if there's no existing delegate
+          // AND no session delegate (session delegates are called for async/await too)
+          if task.delegate == nil, 
+             task.state != .running,
+             (task.value(forKey: "session") as? URLSession)?.delegate == nil {
+            task.delegate = AsyncTaskDelegate(instrumentation: self, sessionTaskId: taskId)
+          }
+          return
+        }
+      }
+
       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
@@ -770,3 +809,32 @@ class FakeDelegate: NSObject, URLSessionTaskDelegate {
   func urlSession(_ session: URLSession, task: URLSessionTask,
                   didCompleteWithError error: Error?) {}
 }
+
+@available(macOS 12.0, iOS 15.0, tvOS 15.0, watchOS 8.0, *)
+class AsyncTaskDelegate: NSObject, URLSessionTaskDelegate {
+  private weak var instrumentation: URLSessionInstrumentation?
+  private let sessionTaskId: String
+  
+  init(instrumentation: URLSessionInstrumentation, sessionTaskId: String) {
+    self.instrumentation = instrumentation
+    self.sessionTaskId = sessionTaskId
+    super.init()
+  }
+  
+  func urlSession(_ session: URLSession, task: URLSessionTask,
+                  didCompleteWithError error: Error?) {
+    guard let instrumentation = instrumentation else { return }
+    
+    // Get the task ID that was set when the task was created
+    let taskId = sessionTaskId
+    
+    if let error = error {
+      let status = (task.response as? HTTPURLResponse)?.statusCode ?? 0
+      URLSessionLogger.logError(error, dataOrFile: nil, statusCode: status,
+                                instrumentation: instrumentation, sessionTaskId: taskId)
+    } else if let response = task.response {
+      URLSessionLogger.logResponse(response, dataOrFile: nil,
+                                   instrumentation: instrumentation, sessionTaskId: taskId)
+    }
+  }
+}

Tests/InstrumentationTests/URLSessionTests/URLSessionInstrumentationTests.swift

@@ -509,6 +509,9 @@ class URLSessionInstrumentationTests: XCTestCase {
     let string = String(decoding: data, as: UTF8.self)
     print(string)
 
+    // Note: These tests were passing incorrectly. The async/await methods
+    // introduced in iOS 15/macOS 12 are NOT instrumented at all, which is what
+    // testAsyncAwaitMethodsAreNotInstrumented demonstrates.
     XCTAssertTrue(URLSessionInstrumentationTests.checker.shouldInstrumentCalled)
     XCTAssertTrue(URLSessionInstrumentationTests.checker.nameSpanCalled)
     XCTAssertTrue(URLSessionInstrumentationTests.checker.spanCustomizationCalled)
@@ -810,4 +813,69 @@ class URLSessionInstrumentationTests: XCTestCase {
 
     try await task.value
   }
+  
+  @available(macOS 12.0, iOS 15.0, tvOS 15.0, watchOS 8.0, *)
+  public func testAsyncAwaitMethodsDoNotCompleteSpans() async throws {
+    let request = URLRequest(url: URL(string: "http://localhost:33333/success")!)
+    
+    // Test data(for:) method - the new async/await API introduced in iOS 15
+    let (data, response) = try await URLSession.shared.data(for: request)
+    
+    guard let httpResponse = response as? HTTPURLResponse else {
+      XCTFail("Response should be HTTPURLResponse")
+      return
+    }
+    
+    XCTAssertEqual(httpResponse.statusCode, 200, "Request should succeed")
+    XCTAssertNotNil(data, "Should receive data")
+
+    XCTAssertTrue(URLSessionInstrumentationTests.checker.receivedResponseCalled, "receivedResponse should be called")
+    XCTAssertNotNil(URLSessionInstrumentationTests.requestCopy?.allHTTPHeaderFields?[W3CTraceContextPropagator.traceparent], "Headers are injected")
+    XCTAssertTrue(URLSessionInstrumentationTests.checker.shouldInstrumentCalled, "shouldInstrument should be called")
+    XCTAssertTrue(URLSessionInstrumentationTests.checker.createdRequestCalled, "createdRequest should be called")
+  }
+  
+  @available(macOS 12.0, iOS 15.0, tvOS 15.0, watchOS 8.0, *)
+  public func testAsyncAwaitDownloadMethodsAreNotInstrumented() async throws {
+    let url = URL(string: "http://localhost:33333/success")!
+    
+    // Test download(from:) method
+    let (fileURL, response) = try await URLSession.shared.download(from: url)
+    
+    guard let httpResponse = response as? HTTPURLResponse else {
+      XCTFail("Response should be HTTPURLResponse")
+      return
+    }
+    
+    XCTAssertEqual(httpResponse.statusCode, 200, "Request should succeed")
+    XCTAssertNotNil(fileURL, "Should receive file URL")
+    
+    XCTAssertTrue(URLSessionInstrumentationTests.checker.shouldInstrumentCalled, "shouldInstrument should be called")
+    XCTAssertTrue(URLSessionInstrumentationTests.checker.createdRequestCalled, "createdRequest should be called")
+    XCTAssertTrue(URLSessionInstrumentationTests.checker.receivedResponseCalled, "receivedResponse should be called")
+    
+    // Clean up downloaded file
+    try? FileManager.default.removeItem(at: fileURL)
+  }
+  
+  @available(macOS 12.0, iOS 15.0, tvOS 15.0, watchOS 8.0, *)
+  public func testAsyncAwaitUploadMethodsAreNotInstrumented() async throws {
+    let url = URL(string: "http://localhost:33333/success")!
+    let request = URLRequest(url: url)
+    
+    // Test upload(for:from:) method
+    let (data, response) = try await URLSession.shared.upload(for: request, from: Data())
+    
+    guard let httpResponse = response as? HTTPURLResponse else {
+      XCTFail("Response should be HTTPURLResponse")
+      return
+    }
+    
+    XCTAssertEqual(httpResponse.statusCode, 200, "Request should succeed")
+    XCTAssertNotNil(data, "Should receive response data")
+
+    XCTAssertTrue(URLSessionInstrumentationTests.checker.shouldInstrumentCalled, "shouldInstrument should be called")
+    XCTAssertTrue(URLSessionInstrumentationTests.checker.createdRequestCalled, "createdRequest should be called")
+    XCTAssertTrue(URLSessionInstrumentationTests.checker.receivedResponseCalled, "receivedResponse should be called")
+  }
 }