Merge pull request #1696 from DataDog/ganeshnj/fix/resource-attributes-provider

ganeshnj · web-flow · commit f1f3d40649b4 · 2024-03-05T18:58:32.000+01:00
fix: pass through data when network request completes
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: \(String(describing: 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,116 @@ 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) {
+        }
+    }
 }
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
@@ -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)
     }
 }
diff --git a/DatadogRUM/Sources/RUMConfiguration.swift b/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?
 

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 {`
Original file line number	Diff line number	Diff line change
`@@ -10,6 +10,9 @@ import XCTest`
`10`	`10`
`11`	`11`	`class URLSessionSwizzlerTests: XCTestCase {`
`12`	`12`	`func testSwizzling_dataTaskWithCompletion() throws {`
	`13`	`+ let didReceive = expectation(description: "didReceive")`
	`14`	`+ didReceive.expectedFulfillmentCount = 2`
	`15`	`+`
`13`	`16`	`let didInterceptCompletion = expectation(description: "interceptCompletion")`
`14`	`17`	`didInterceptCompletion.expectedFulfillmentCount = 2`
`15`	`18`
`@@ -18,6 +21,8 @@ class URLSessionSwizzlerTests: XCTestCase {`
`18`	`21`	`try swizzler.swizzle(`
`19`	`22`	`interceptCompletionHandler: { _, _, _ in`
`20`	`23`	`didInterceptCompletion.fulfill()`
	`24`	`+ }, didReceive: { _, _ in`
	`25`	`+ didReceive.fulfill()`
`21`	`26`	`}`
`22`	`27`	`)`
`23`	`28`
`@@ -30,6 +35,6 @@ class URLSessionSwizzlerTests: XCTestCase {`
`30`	`35`	`session.dataTask(with: url) { _, _, _ in }.resume() // not intercepted`
`31`	`36`	`session.dataTask(with: URLRequest(url: url)) { _, _, _ in }.resume() // not intercepted`
`32`	`37`
`33`		`- wait(for: [didInterceptCompletion], timeout: 5)`
	`38`	`+ wait(for: [didReceive, didInterceptCompletion], timeout: 5)`
`34`	`39`	`}`
`35`	`40`	`}`