From 50bf3f5b68dbe4a2ed2b889e6106a9db28744870 Mon Sep 17 00:00:00 2001 From: Ganesh Jangir Date: Fri, 23 Feb 2024 17:22:43 +0100 Subject: [PATCH 1/3] fix: pass through data when network request completes --- Datadog/Example/ExampleAppDelegate.swift | 6 +- ...tworkInstrumentationIntegrationTests.swift | 116 +++++++++++++++++- .../RUM/StartingRUMSessionTests.swift | 8 +- .../NetworkInstrumentationFeature.swift | 4 +- .../DatadogURLSessionDelegate.swift | 1 + .../NetworkInstrumentationSwizzler.swift | 8 +- .../URLSession/URLSessionSwizzler.swift | 31 ++++- .../URLSession/URLSessionTask+Tracking.swift | 16 +++ .../URLSessionSwizzlerTests.swift | 7 +- DatadogRUM/Sources/RUMConfiguration.swift | 2 + 10 files changed, 183 insertions(+), 16 deletions(-) diff --git a/Datadog/Example/ExampleAppDelegate.swift b/Datadog/Example/ExampleAppDelegate.swift index c86121c55d..e0785c6b68 100644 --- 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 index b7108346db..553c048001 100644 --- 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,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) + } + } } diff --git a/Datadog/IntegrationUnitTests/RUM/StartingRUMSessionTests.swift b/Datadog/IntegrationUnitTests/RUM/StartingRUMSessionTests.swift index 8dc3a52fc8..9eeb75343d 100644 --- a/Datadog/IntegrationUnitTests/RUM/StartingRUMSessionTests.swift +++ b/Datadog/IntegrationUnitTests/RUM/StartingRUMSessionTests.swift @@ -143,14 +143,14 @@ 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()) @@ -158,7 +158,7 @@ class StartingRUMSessionTests: XCTestCase { 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 index 28f7d97d8d..3b8e451db0 100644 --- 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 index 1462812d0b..ac1e33c77f 100644 --- 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 index f694655db5..8c6c9d6dbf 100644 --- 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 index 2143952fee..5e9450fa8e 100644 --- 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,6 +151,10 @@ 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 @@ -144,6 +162,7 @@ internal final class URLSessionSwizzler { } } _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 index 6c3f19035b..da1a44f42c 100644 --- 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 index 75671d8f53..cfb515f7ae 100644 --- 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 index e1c4824ca3..b7bb27d563 100644 --- 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? From 3daf624295382baa4f6c5aa9ec8d3da05733ee3f Mon Sep 17 00:00:00 2001 From: Ganesh Jangir Date: Sat, 24 Feb 2024 23:44:58 +0100 Subject: [PATCH 2/3] fix warning --- Datadog/Example/ExampleAppDelegate.swift | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Datadog/Example/ExampleAppDelegate.swift b/Datadog/Example/ExampleAppDelegate.swift index e0785c6b68..45f764ed18 100644 --- a/Datadog/Example/ExampleAppDelegate.swift +++ b/Datadog/Example/ExampleAppDelegate.swift @@ -68,7 +68,7 @@ class ExampleAppDelegate: UIResponder, UIApplicationDelegate { applicationID: Environment.readRUMApplicationID(), urlSessionTracking: .init( resourceAttributesProvider: { req, resp, data, err in - print("⭐️ [Attributes Provider] data: \(data)") + print("⭐️ [Attributes Provider] data: \(String(describing: data))") return [:] }), trackBackgroundEvents: true, From 5cf5f8c84fdd5f238a11c87acb1c4a4b72c8eeca Mon Sep 17 00:00:00 2001 From: Ganesh Jangir Date: Mon, 4 Mar 2024 13:19:12 +0100 Subject: [PATCH 3/3] Update Datadog/IntegrationUnitTests/Public/NetworkInstrumentationIntegrationTests.swift Co-authored-by: Maciek Grzybowski --- .../Public/NetworkInstrumentationIntegrationTests.swift | 1 - 1 file changed, 1 deletion(-) diff --git a/Datadog/IntegrationUnitTests/Public/NetworkInstrumentationIntegrationTests.swift b/Datadog/IntegrationUnitTests/Public/NetworkInstrumentationIntegrationTests.swift index 553c048001..209309cf9a 100644 --- a/Datadog/IntegrationUnitTests/Public/NetworkInstrumentationIntegrationTests.swift +++ b/Datadog/IntegrationUnitTests/Public/NetworkInstrumentationIntegrationTests.swift @@ -199,7 +199,6 @@ class NetworkInstrumentationIntegrationTests: XCTestCase { class InstrumentedSessionDelegate: NSObject, URLSessionDataDelegate { func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive data: Data) { - print(data) } } }