fix: pass through data when network request completes

Author: ganeshnj

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: \(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,115 @@ 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
+        )
+        var request = URLRequest(url: URL(string: "https://www.datadoghq.com/")!)
+        request.httpMethod = "GET"
+
+        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 {
+    }
 }

DatadogInternal/Sources/NetworkInstrumentation/NetworkInstrumentationFeature.swift

@@ -91,15 +91,9 @@ internal final class NetworkInstrumentationFeature: DatadogFeature {
             delegateClass: configuration.delegateClass,
             interceptDidFinishCollecting: { [weak self] session, task, metrics in
                 self?.task(task, didFinishCollecting: metrics)
-
-                if #available(iOS 15, tvOS 15, *) {
-                    // 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)
-                }
             },
             interceptDidCompleteWithError: { [weak self] session, task, error in
-                self?.task(task, didCompleteWithError: error)
+                self?.task(task, data: nil, didCompleteWithError: error)
             }
         )
 
@@ -111,8 +105,8 @@ internal final class NetworkInstrumentationFeature: DatadogFeature {
         )
 
         try swizzler.swizzle(
-            interceptCompletionHandler: { [weak self] task, _, error in
-                self?.task(task, didCompleteWithError: error)
+            interceptCompletionHandler: { [weak self] task, data, error in
+                self?.task(task, data: data, didCompleteWithError: error)
             }
         )
     }
@@ -231,12 +225,16 @@ extension NetworkInstrumentationFeature {
     /// - Parameters:
     ///   - task: The task that has finished transferring data.
     ///   - error: If an error occurred, an error object indicating how the transfer failed, otherwise NULL.
-    func task(_ task: URLSessionTask, didCompleteWithError error: Error?) {
-        queue.async { [weak self] in
+    func task(_ task: URLSessionTask, data: Data?, didCompleteWithError error: Error?) {
+        queue.async { [weak self, data] in
             guard let self = self, let interception = self.interceptions[task] else {
                 return
             }
 
+            if let data = data {
+                self.interceptions[task]?.register(nextData: data)
+            }
+
             interception.register(
                 response: task.response,
                 error: error

DatadogInternal/Sources/NetworkInstrumentation/URLSession/URLSessionInterceptor.swift

@@ -70,6 +70,6 @@ public struct URLSessionInterceptor {
     ///   - task: The task that has finished transferring data.
     ///   - error: If an error occurred, an error object indicating how the transfer failed, otherwise NULL.
     public func task(_ task: URLSessionTask, didCompleteWithError error: Error?) {
-        feature.task(task, didCompleteWithError: error)
+        feature.task(task, data: nil, didCompleteWithError: error)
     }
 }

DatadogInternal/Tests/NetworkInstrumentation/NetworkInstrumentationFeatureTests.swift

@@ -703,7 +703,7 @@ class NetworkInstrumentationFeatureTests: XCTestCase {
                 { feature.intercept(task: tasks.randomElement()!, additionalFirstPartyHosts: nil) },
                 { feature.task(tasks.randomElement()!, didReceive: .mockRandom()) },
                 { feature.task(tasks.randomElement()!, didFinishCollecting: .mockAny()) },
-                { feature.task(tasks.randomElement()!, didCompleteWithError: nil) },
+                { feature.task(tasks.randomElement()!, data: nil, didCompleteWithError: nil) },
                 { try? feature.bind(configuration: .init(delegateClass: MockDelegate.self)) },
                 { feature.unbind(delegateClass: MockDelegate.self) }
             ],