fix: pass through data when network request completes

ganeshnj · ganeshnj · commit 7e761a82e687 · 2024-02-23T17:22:43.000+01:00
diff --git a/Datadog/Example/ExampleAppDelegate.swift b/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
diff --git a/Datadog/IntegrationUnitTests/Public/NetworkInstrumentationIntegrationTests.swift b/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,172 @@ 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)
+    }
+
+    @available(iOS 13.0, tvOS 13.0, *)
+    func testResourceAttributesProvider_givenURLSessionDataWithRequest() async throws {
+        guard #available(iOS 16, tvOS 16, *) else {
+            return
+        }
+        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
+                        XCTAssertNil(data)
+                        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 (data, _) = try await URLSession.shared.data(for: request, delegate: InstrumentedSessionDelegate())
+        let string = String(decoding: data, as: UTF8.self)
+        print(string)
+
+        var taskDataCount = 0
+        let (taskData, _) = try await session.data(for: request)
+
+        wait(for: [providerExpectation], timeout: 10)
+        XCTAssertEqual(providerDataCount, taskDataCount)
+        XCTAssertEqual(providerData, taskData)
+    }
+
+    class InstrumentedSessionDelegate: NSObject, URLSessionDataDelegate {
+        func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive data: Data) {
+            print(data)
+        }
+    }
 }
diff --git a/Datadog/IntegrationUnitTests/RUM/StartingRUMSessionTests.swift b/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")
diff --git a/DatadogInternal/Sources/NetworkInstrumentation/NetworkInstrumentationFeature.swift b/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)
             }
         )
     }
diff --git a/DatadogInternal/Sources/NetworkInstrumentation/URLSession/DatadogURLSessionDelegate.swift b/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 {
diff --git a/DatadogInternal/Sources/NetworkInstrumentation/URLSession/NetworkInstrumentationSwizzler.swift b/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.
diff --git a/DatadogInternal/Sources/NetworkInstrumentation/URLSession/URLSessionSwizzler.swift b/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
                 }
             }
diff --git a/DatadogInternal/Sources/NetworkInstrumentation/URLSession/URLSessionTask+Tracking.swift b/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?
diff --git a/DatadogInternal/Tests/NetworkInstrumentation/URLSessionSwizzlerTests.swift b/DatadogInternal/Tests/NetworkInstrumentation/URLSessionSwizzlerTests.swift

Original file line number	Diff line number	Diff line change
`@@ -92,7 +92,7 @@ internal final class NetworkInstrumentationFeature: DatadogFeature {`
`92`	`92`	`interceptDidFinishCollecting: { [weak self] session, task, metrics in`
`93`	`93`	`self?.task(task, didFinishCollecting: metrics)`
`94`	`94`
`95`		`- if #available(iOS 15, tvOS 15, *) {`
	`95`	`+ if #available(iOS 15, tvOS 15, *), !task.dd.hasCompletion {`
`96`	`96`	`// iOS 15 and above, didCompleteWithError is not called hence we use task state to detect task completion`
`97`	`97`	`// while prior to iOS 15, task state doesn't change to completed hence we use didCompleteWithError to detect task completion`
`98`	`98`	`self?.task(task, didCompleteWithError: task.error)`
`@@ -113,6 +113,8 @@ internal final class NetworkInstrumentationFeature: DatadogFeature {`
`113`	`113`	`try swizzler.swizzle(`
`114`	`114`	`interceptCompletionHandler: { [weak self] task, _, error in`
`115`	`115`	`self?.task(task, didCompleteWithError: error)`
	`116`	`+ }, didReceive: { [weak self] task, data in`
	`117`	`+ self?.task(task, didReceive: data)`
`116`	`118`	`}`
`117`	`119`	`)`
`118`	`120`	`}`
Original file line number	Diff line number	Diff line change
`@@ -139,6 +139,7 @@ open class DatadogURLSessionDelegate: NSObject, URLSessionDataDelegate {`
`139`	`139`	`try swizzler.swizzle(`
`140`	`140`	`interceptCompletionHandler: { [weak self] task, _, error in`
`141`	`141`	`self?.interceptor?.task(task, didCompleteWithError: error)`
	`142`	`+ }, didReceive: { _, _ in`
`142`	`143`	`}`
`143`	`144`	`)`
`144`	`145`	`} catch {`