fix: pass through data when network request completes

Datadog/Example/ExampleAppDelegate.swift

@@ -66,7 +66,11 @@ class ExampleAppDelegate: UIResponder, UIApplicationDelegate {
         RUM.enable(
             with: RUM.Configuration(
                 applicationID: Environment.readRUMApplicationID(),
-                urlSessionTracking: .init(firstPartyHostsTracing: .trace(hosts: [], sampleRate: 100)),
+                urlSessionTracking: .init(
+                    resourceAttributesProvider: { req, resp, data, err in
+                        print("⭐️ [Attributes Provider] data: \(String(describing: data))")
+                        return [:]
+            }),
                 trackBackgroundEvents: true,
                 customEndpoint: Environment.readCustomRUMURL(),
                 telemetrySampleRate: 100

Datadog/IntegrationUnitTests/Public/NetworkInstrumentationIntegrationTests.swift

@@ -8,6 +8,7 @@ import XCTest
 import TestUtilities
 import DatadogInternal
 
+@testable import DatadogRUM
 @testable import DatadogTrace
 @testable import DatadogCore
 
@@ -50,7 +51,7 @@ class NetworkInstrumentationIntegrationTests: XCTestCase {
         core.flushAndTearDown()
         core = nil
     }
-    
+
     func testParentSpanPropagation() throws {
         let expectation = expectation(description: "request completes")
         // Given
@@ -88,4 +89,117 @@ class NetworkInstrumentationIntegrationTests: XCTestCase {
 
     class MockDelegate: NSObject, URLSessionDataDelegate {
     }
+
+    func testResourceAttributesProvider_givenURLSessionDataTaskRequest() {
+        core = DatadogCoreProxy(
+            context: .mockWith(
+                env: "test",
+                version: "1.1.1",
+                serverTimeOffset: 123
+            )
+        )
+
+        let providerExpectation = expectation(description: "provider called")
+        var providerDataCount = 0
+        RUM.enable(
+            with: .init(
+                applicationID: .mockAny(),
+                urlSessionTracking: .init(
+                    resourceAttributesProvider: { req, resp, data, err in
+                        XCTAssertNotNil(data)
+                        XCTAssertTrue(data!.count > 0)
+                        providerDataCount = data!.count
+                        providerExpectation.fulfill()
+                        return [:]
+                })
+            ),
+            in: core
+        )
+
+        URLSessionInstrumentation.enable(
+            with: .init(
+                delegateClass: InstrumentedSessionDelegate.self
+            ),
+            in: core
+        )
+
+        let session = URLSession(
+            configuration: .ephemeral,
+            delegate: InstrumentedSessionDelegate(),
+            delegateQueue: nil
+        )
+        var request = URLRequest(url: URL(string: "https://www.datadoghq.com/")!)
+        request.httpMethod = "GET"
+
+        let task = session.dataTask(with: request)
+        task.resume()
+
+        wait(for: [providerExpectation], timeout: 10)
+        XCTAssertTrue(providerDataCount > 0)
+    }
+
+    func testResourceAttributesProvider_givenURLSessionDataTaskRequestWithCompletionHandler() {
+        core = DatadogCoreProxy(
+            context: .mockWith(
+                env: "test",
+                version: "1.1.1",
+                serverTimeOffset: 123
+            )
+        )
+
+        let providerExpectation = expectation(description: "provider called")
+        var providerDataCount = 0
+        var providerData: Data?
+        RUM.enable(
+            with: .init(
+                applicationID: .mockAny(),
+                urlSessionTracking: .init(
+                    resourceAttributesProvider: { req, resp, data, err in
+                        XCTAssertNotNil(data)
+                        XCTAssertTrue(data!.count > 0)
+                        providerDataCount = data!.count
+                        data.map { providerData = $0 }
+                        providerExpectation.fulfill()
+                        return [:]
+                })
+            ),
+            in: core
+        )
+
+        URLSessionInstrumentation.enable(
+            with: .init(
+                delegateClass: InstrumentedSessionDelegate.self
+            ),
+            in: core
+        )
+
+        let session = URLSession(
+            configuration: .ephemeral,
+            delegate: InstrumentedSessionDelegate(),
+            delegateQueue: nil
+        )
+        let request = URLRequest(url: URL(string: "https://www.datadoghq.com/")!)
+
+        let taskExpectation = self.expectation(description: "task completed")
+        var taskDataCount = 0
+        var taskData: Data?
+        let task = session.dataTask(with: request) { data, _, _ in
+            XCTAssertNotNil(data)
+            XCTAssertTrue(data!.count > 0)
+            taskDataCount = data!.count
+            data.map { taskData = $0 }
+            taskExpectation.fulfill()
+        }
+        task.resume()
+
+        wait(for: [providerExpectation, taskExpectation], timeout: 10)
+        XCTAssertEqual(providerDataCount, taskDataCount)
+        XCTAssertEqual(providerData, taskData)
+    }
+
+    class InstrumentedSessionDelegate: NSObject, URLSessionDataDelegate {
+        func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive data: Data) {
+            print(data)
+        }
+    }
 }

Datadog/IntegrationUnitTests/RUM/StartingRUMSessionTests.swift

@@ -143,22 +143,22 @@ class StartingRUMSessionTests: XCTestCase {
         let rumTime = DateProviderMock()
         rumConfig.dateProvider = rumTime
         rumConfig.trackBackgroundEvents = .mockRandom() // no matter BET state
-        
+
         // When
         rumTime.now = sdkInitTime
         RUM.enable(with: rumConfig, in: core)
-        
+
         rumTime.now = firstRUMTime
         RUMMonitor.shared(in: core).startView(key: "key", name: "FirstView")
-        
+
         // Then
         let session = try RUMSessionMatcher
             .groupMatchersBySessions(try core.waitAndReturnRUMEventMatchers())
             .takeSingle()
 
         XCTAssertEqual(session.views.count, 1)
         XCTAssertTrue(try session.has(sessionPrecondition: .backgroundLaunch), "Session must be marked as 'background launch'")
-        
+
         let firstView = try XCTUnwrap(session.views.first)
         XCTAssertFalse(firstView.isApplicationLaunchView(), "Session should not begin with 'app launch' view")
         XCTAssertEqual(firstView.name, "FirstView")

DatadogInternal/Sources/NetworkInstrumentation/NetworkInstrumentationFeature.swift

@@ -92,7 +92,7 @@ internal final class NetworkInstrumentationFeature: DatadogFeature {
             interceptDidFinishCollecting: { [weak self] session, task, metrics in
                 self?.task(task, didFinishCollecting: metrics)
 
-                if #available(iOS 15, tvOS 15, *) {
+                if #available(iOS 15, tvOS 15, *), !task.dd.hasCompletion {
                     // iOS 15 and above, didCompleteWithError is not called hence we use task state to detect task completion
                     // while prior to iOS 15, task state doesn't change to completed hence we use didCompleteWithError to detect task completion
                     self?.task(task, didCompleteWithError: task.error)
@@ -113,6 +113,8 @@ internal final class NetworkInstrumentationFeature: DatadogFeature {
         try swizzler.swizzle(
             interceptCompletionHandler: { [weak self] task, _, error in
                 self?.task(task, didCompleteWithError: error)
+            }, didReceive: { [weak self] task, data in
+                self?.task(task, didReceive: data)
             }
         )
     }

DatadogInternal/Sources/NetworkInstrumentation/URLSession/DatadogURLSessionDelegate.swift

@@ -139,6 +139,7 @@ open class DatadogURLSessionDelegate: NSObject, URLSessionDataDelegate {
             try swizzler.swizzle(
                 interceptCompletionHandler: { [weak self] task, _, error in
                     self?.interceptor?.task(task, didCompleteWithError: error)
+                }, didReceive: { _, _ in
                 }
             )
         } catch {

DatadogInternal/Sources/NetworkInstrumentation/URLSession/NetworkInstrumentationSwizzler.swift

@@ -23,9 +23,13 @@ internal final class NetworkInstrumentationSwizzler {
 
     /// Swizzles `URLSession.dataTask(with:completionHandler:)` methods (with `URL` and `URLRequest`).
     func swizzle(
-        interceptCompletionHandler: @escaping (URLSessionTask, Data?, Error?) -> Void
+        interceptCompletionHandler: @escaping (URLSessionTask, Data?, Error?) -> Void,
+        didReceive: @escaping (URLSessionTask, Data) -> Void
     ) throws {
-        try urlSessionSwizzler.swizzle(interceptCompletionHandler: interceptCompletionHandler)
+        try urlSessionSwizzler.swizzle(
+            interceptCompletionHandler: interceptCompletionHandler,
+            didReceive: didReceive
+        )
     }
 
     /// Swizzles `URLSessionTask.resume()` method.

DatadogInternal/Sources/NetworkInstrumentation/URLSession/URLSessionSwizzler.swift

@@ -18,19 +18,26 @@ internal final class URLSessionSwizzler {
 
     /// Swizzles `URLSession.dataTask(with:completionHandler:)` methods (with `URL` and `URLRequest`).
     func swizzle(
-        interceptCompletionHandler: @escaping (URLSessionTask, Data?, Error?) -> Void
+        interceptCompletionHandler: @escaping (URLSessionTask, Data?, Error?) -> Void,
+        didReceive: @escaping (URLSessionTask, Data) -> Void
     ) throws {
         lock.lock()
         defer { lock.unlock() }
         dataTaskURLRequestCompletionHandler = try DataTaskURLRequestCompletionHandler.build()
-        dataTaskURLRequestCompletionHandler?.swizzle(interceptCompletion: interceptCompletionHandler)
+        dataTaskURLRequestCompletionHandler?.swizzle(
+            interceptCompletion: interceptCompletionHandler,
+            didReceive: didReceive
+        )
 
         if #available(iOS 13.0, *) {
             // Prior to iOS 13.0 the `URLSession.dataTask(with:url, completionHandler:handler)` makes an internal
             // call to `URLSession.dataTask(with:request, completionHandler:handler)`. To avoid duplicated call
             // to the callback, we don't apply below swizzling prior to iOS 13.
             dataTaskURLCompletionHandler = try DataTaskURLCompletionHandler.build()
-            dataTaskURLCompletionHandler?.swizzle(interceptCompletion: interceptCompletionHandler)
+            dataTaskURLCompletionHandler?.swizzle(
+                interceptCompletion: interceptCompletionHandler,
+                didReceive: didReceive
+            )
         }
     }
 
@@ -71,7 +78,8 @@ internal final class URLSessionSwizzler {
         }
 
         func swizzle(
-            interceptCompletion: @escaping (URLSessionTask, Data?, Error?) -> Void
+            interceptCompletion: @escaping (URLSessionTask, Data?, Error?) -> Void,
+            didReceive: @escaping (URLSessionTask, Data) -> Void
         ) {
             typealias Signature = @convention(block) (URLSession, URLRequest, CompletionHandler?) -> URLSessionDataTask
             swizzle(method) { previousImplementation -> Signature in
@@ -87,13 +95,18 @@ internal final class URLSessionSwizzler {
 
                     var _task: URLSessionDataTask?
                     let task = previousImplementation(session, Self.selector, request) { data, response, error in
-                        completionHandler(data, response, error)
+                        if let task = _task, let data = data {
+                            didReceive(task, data)
+                        }
 
                         if let task = _task { // sanity check, should always succeed
                             interceptCompletion(task, data, error)
                         }
+
+                        completionHandler(data, response, error)
                     }
                     _task = task
+                    _task?.dd.hasCompletion = true
                     return task
                 }
             }
@@ -121,7 +134,8 @@ internal final class URLSessionSwizzler {
         }
 
         func swizzle(
-            interceptCompletion: @escaping (URLSessionTask, Data?, Error?) -> Void
+            interceptCompletion: @escaping (URLSessionTask, Data?, Error?) -> Void,
+            didReceive: @escaping (URLSessionTask, Data) -> Void
         ) {
             typealias Signature = @convention(block) (URLSession, URL, CompletionHandler?) -> URLSessionDataTask
             swizzle(method) { previousImplementation -> Signature in
@@ -137,13 +151,18 @@ internal final class URLSessionSwizzler {
 
                     var _task: URLSessionDataTask?
                     let task = previousImplementation(session, Self.selector, url) { data, response, error in
+                        if let task = _task, let data = data {
+                            didReceive(task, data)
+                        }
+
                         completionHandler(data, response, error)
 
                         if let task = _task { // sanity check, should always succeed
                             interceptCompletion(task, data, error)
                         }
                     }
                     _task = task
+                    _task?.dd.hasCompletion = true
                     return task
                 }
             }

DatadogInternal/Sources/NetworkInstrumentation/URLSession/URLSessionTask+Tracking.swift

@@ -33,4 +33,20 @@ extension DatadogExtension where ExtendedType: URLSessionTask {
 
         return session.delegate
     }
+
+    var hasCompletion: Bool {
+        get {
+            let value = objc_getAssociatedObject(type, &hasCompletionKey) as? Bool
+            return value == true
+        }
+        set {
+            if newValue {
+                objc_setAssociatedObject(type, &hasCompletionKey, true, .OBJC_ASSOCIATION_RETAIN_NONATOMIC)
+            } else {
+                objc_setAssociatedObject(type, &hasCompletionKey, nil, .OBJC_ASSOCIATION_RETAIN_NONATOMIC)
+            }
+        }
+    }
 }
+
+private var hasCompletionKey: Void?

DatadogInternal/Tests/NetworkInstrumentation/URLSessionSwizzlerTests.swift

@@ -10,6 +10,9 @@ import XCTest
 
 class URLSessionSwizzlerTests: XCTestCase {
     func testSwizzling_dataTaskWithCompletion() throws {
+        let didReceive = expectation(description: "didReceive")
+        didReceive.expectedFulfillmentCount = 2
+
         let didInterceptCompletion = expectation(description: "interceptCompletion")
         didInterceptCompletion.expectedFulfillmentCount = 2
 
@@ -18,6 +21,8 @@ class URLSessionSwizzlerTests: XCTestCase {
         try swizzler.swizzle(
             interceptCompletionHandler: { _, _, _ in
                 didInterceptCompletion.fulfill()
+            }, didReceive: { _, _ in
+                didReceive.fulfill()
             }
         )
 
@@ -30,6 +35,6 @@ class URLSessionSwizzlerTests: XCTestCase {
         session.dataTask(with: url) { _, _, _ in }.resume() // not intercepted
         session.dataTask(with: URLRequest(url: url)) { _, _, _ in }.resume() // not intercepted
 
-        wait(for: [didInterceptCompletion], timeout: 5)
+        wait(for: [didReceive, didInterceptCompletion], timeout: 5)
     }
 }

DatadogRUM/Sources/RUMConfiguration.swift

@@ -221,6 +221,8 @@ extension RUM {
             /// attributes for RUM resource based on the provided request, response, data, and error.
             /// Keep the implementation fast and do not make any assumptions on the thread used to run it.
             ///
+            /// Note: This is not supported for async-await APIs.
+            ///
             /// Default: `nil`.
             public var resourceAttributesProvider: RUM.ResourceAttributesProvider?