Merge pull request #473 from DataDog/ncreated/RUMM-1154-user-can-add-…

…attributes-to-auto-instrumented-rum-resources

RUMM-1154 User can add attributes to auto instrumented RUM Resources

Datadog/Example/Scenarios/RUM/RUMScenarios.swift

@@ -129,6 +129,7 @@ final class RUMURLSessionResourcesScenario: URLSessionBaseScenario, TestScenario
     override func configureSDK(builder: Datadog.Configuration.Builder) {
         _ = builder
             .trackUIKitRUMViews()
+            .setRUMResourceAttributesProvider(rumResourceAttributesProvider(request:response:data:error:))
 
         super.configureSDK(builder: builder) // applies the `trackURLSession(firstPartyHosts:)`
     }
@@ -142,6 +143,7 @@ final class RUMNSURLSessionResourcesScenario: URLSessionBaseScenario, TestScenar
     override func configureSDK(builder: Datadog.Configuration.Builder) {
         _ = builder
             .trackUIKitRUMViews()
+            .setRUMResourceAttributesProvider(rumResourceAttributesProvider(request:response:data:error:))
 
         super.configureSDK(builder: builder) // applies the `trackURLSession(firstPartyHosts:)`
     }
@@ -203,3 +205,41 @@ final class RUMScrubbingScenario: TestScenario {
             }
     }
 }
+
+// MARK: - Helpers
+
+private func rumResourceAttributesProvider(
+    request: URLRequest,
+    response: URLResponse?,
+    data: Data?,
+    error: Error?
+) -> [AttributeKey: AttributeValue]? {
+    /// Apples new-line separated text format to  response headers.
+    func format(headers: [AnyHashable: Any]) -> String {
+        var formattedHeaders: [String] = []
+        headers.forEach { key, value in
+            formattedHeaders.append("\(String(describing: key)): \(String(describing: value))")
+        }
+        return formattedHeaders.joined(separator: "\n")
+    }
+
+    var responseBodyValue: String?
+    var responseHeadersValue: String?
+    var errorDetailsValue: String?
+
+    if let responseHeaders = (response as? HTTPURLResponse)?.allHeaderFields {
+        responseHeadersValue = format(headers: responseHeaders)
+    }
+    if let data = data {
+        responseBodyValue = String(data: data, encoding: .utf8) ?? "<not an UTF-8 data>"
+    }
+    if let error = error {
+        errorDetailsValue = String(describing: error)
+    }
+
+    return [
+        "response.body.truncated" : responseBodyValue.flatMap { String($0.prefix(128)) },
+        "response.headers" : responseHeadersValue,
+        "response.error" : errorDetailsValue,
+    ]
+}

Sources/Datadog/Core/FeaturesConfiguration.swift

@@ -54,6 +54,8 @@ internal struct FeaturesConfiguration {
         let userDefinedFirstPartyHosts: Set<String>
         /// URLs used internally by the SDK - they are not instrumented.
         let sdkInternalURLs: Set<String>
+        /// An optional RUM Resource attributes provider.
+        let rumAttributesProvider: URLSessionRUMAttributesProvider?
 
         /// If the Tracing instrumentation should be enabled.
         let instrumentTracing: Bool
@@ -215,6 +217,7 @@ extension FeaturesConfiguration {
                         tracesEndpoint.url,
                         rumEndpoint.url
                     ],
+                    rumAttributesProvider: configuration.rumResourceAttributesProvider,
                     instrumentTracing: configuration.tracingEnabled,
                     instrumentRUM: configuration.rumEnabled
                 )
@@ -227,6 +230,14 @@ extension FeaturesConfiguration {
                 )
                 consolePrint("\(error)")
             }
+        } else if configuration.rumResourceAttributesProvider != nil {
+            let error = ProgrammerError(
+                description: """
+                To use `.setRUMResourceAttributesProvider(_:)` URLSession tracking must be enabled
+                with `.trackURLSession(firstPartyHosts:)`.
+                """
+            )
+            consolePrint("\(error)")
         }
 
         if let crashReportingPlugin = configuration.crashReportingPlugin {

Sources/Datadog/DatadogConfiguration.swift

@@ -178,6 +178,7 @@ extension Datadog {
         private(set) var rumResourceEventMapper: RUMResourceEventMapper?
         private(set) var rumActionEventMapper: RUMActionEventMapper?
         private(set) var rumErrorEventMapper: RUMErrorEventMapper?
+        private(set) var rumResourceAttributesProvider: URLSessionRUMAttributesProvider?
         private(set) var batchSize: BatchSize
         private(set) var uploadFrequency: UploadFrequency
         private(set) var additionalConfiguration: [String: Any]
@@ -247,6 +248,7 @@ extension Datadog {
                     rumResourceEventMapper: nil,
                     rumActionEventMapper: nil,
                     rumErrorEventMapper: nil,
+                    rumResourceAttributesProvider: nil,
                     batchSize: .medium,
                     uploadFrequency: .average,
                     additionalConfiguration: [:],
@@ -500,6 +502,19 @@ extension Datadog {
                 return self
             }
 
+            /// Sets a closure to provide custom attributes for intercepted RUM Resources.
+            ///
+            /// The `provider` closure is called for each `URLSession` task intercepted by the SDK (each automatically collected RUM Resource).
+            /// The closure is called with session task information (`URLRequest`, `URLResponse?`, `Data?` and `Error?`) that can be used to identify the task, inspect its
+            /// values and return custom attributes for the RUM Resource.
+            ///
+            /// - Parameter provider: the closure called for each RUM Resource collected by the SDK. This closure is called with task information and may return custom attributes
+            ///                       for the RUM Resource or `nil` if no attributes should be attached.
+            public func setRUMResourceAttributesProvider(_ provider: @escaping (URLRequest, URLResponse?, Data?, Error?) -> [AttributeKey: AttributeValue]?) -> Builder {
+                configuration.rumResourceAttributesProvider = provider
+                return self
+            }
+
             // MARK: - Crash Reporting Configuration
 
             /// Enables the crash reporting feature.

Sources/Datadog/RUM/AutoInstrumentation/Resources/URLSessionRUMResourcesHandler.swift

@@ -6,13 +6,19 @@
 
 import Foundation
 
+internal typealias URLSessionRUMAttributesProvider = (URLRequest, URLResponse?, Data?, Error?) -> [AttributeKey: AttributeValue]?
+
 internal class URLSessionRUMResourcesHandler: URLSessionInterceptionHandler {
     private let dateProvider: DateProvider
+    /// Attributes-providing callback.
+    /// It is configured by the user and should be used to associate additional RUM attributes with intercepted RUM Resource.
+    let rumAttributesProvider: (URLSessionRUMAttributesProvider)?
 
     // MARK: - Initialization
 
-    init(dateProvider: DateProvider) {
+    init(dateProvider: DateProvider, rumAttributesProvider: (URLSessionRUMAttributesProvider)?) {
         self.dateProvider = dateProvider
+        self.rumAttributesProvider = rumAttributesProvider
     }
 
     // MARK: - Internal
@@ -57,6 +63,14 @@ internal class URLSessionRUMResourcesHandler: URLSessionInterceptionHandler {
             )
         }
 
+        // Get RUM Resource attributes from the user.
+        let userAttributes = rumAttributesProvider?(
+            interception.request,
+            interception.completion?.httpResponse,
+            interception.data,
+            interception.completion?.error
+        ) ?? [:]
+
         if let resourceMetrics = interception.metrics {
             subscriber?.process(
                 command: RUMAddResourceMetricsCommand(
@@ -73,7 +87,7 @@ internal class URLSessionRUMResourcesHandler: URLSessionInterceptionHandler {
                 command: RUMStopResourceCommand(
                     resourceKey: interception.identifier.uuidString,
                     time: dateProvider.currentDate(),
-                    attributes: [:],
+                    attributes: userAttributes,
                     kind: RUMResourceType(response: httpResponse),
                     httpStatusCode: httpResponse.statusCode,
                     size: interception.metrics?.responseSize
@@ -89,7 +103,7 @@ internal class URLSessionRUMResourcesHandler: URLSessionInterceptionHandler {
                     error: error,
                     source: .network,
                     httpStatusCode: interception.completion?.httpResponse?.statusCode,
-                    attributes: [:]
+                    attributes: userAttributes
                 )
             )
         }

Sources/Datadog/URLSessionAutoInstrumentation/DDURLSessionDelegate.swift

@@ -12,8 +12,9 @@ import Foundation
 public protocol __URLSessionDelegateProviding: URLSessionDelegate {
     /// Datadog delegate object.
     /// The class implementing `DDURLSessionDelegateProviding` must ensure that following method calls are forwarded to `ddURLSessionDelegate`:
-    // - `func urlSession(_:task:didFinishCollecting:)`
-    // - `func urlSession(_:task:didCompleteWithError:)`
+    /// - `func urlSession(_:task:didFinishCollecting:)`
+    /// - `func urlSession(_:task:didCompleteWithError:)`
+    /// - `func urlSession(_:dataTask:didReceive:)`
     var ddURLSessionDelegate: DDURLSessionDelegate { get }
 }
 
@@ -22,7 +23,7 @@ public protocol __URLSessionDelegateProviding: URLSessionDelegate {
 ///
 /// All requests made with the `URLSession` instrumented with this delegate will be intercepted by the SDK.
 @objc
-open class DDURLSessionDelegate: NSObject, URLSessionTaskDelegate, __URLSessionDelegateProviding {
+open class DDURLSessionDelegate: NSObject, URLSessionTaskDelegate, URLSessionDataDelegate, __URLSessionDelegateProviding {
     public var ddURLSessionDelegate: DDURLSessionDelegate {
         return self
     }
@@ -76,4 +77,10 @@ open class DDURLSessionDelegate: NSObject, URLSessionTaskDelegate, __URLSessionD
 
         interceptor?.taskCompleted(task: task, error: error)
     }
+
+    public func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive data: Data) {
+        // NOTE: This delegate method is only called for `URLSessionTasks` created without the completion handler.
+
+        interceptor?.taskReceivedData(task: dataTask, data: data)
+    }
 }

Sources/Datadog/URLSessionAutoInstrumentation/Interception/TaskInterception.swift

@@ -16,6 +16,8 @@ internal class TaskInterception {
     internal let isFirstPartyRequest: Bool
     /// Task metrics collected during this interception.
     private(set) var metrics: ResourceMetrics?
+    /// Task data received during this interception. Can be `nil` if task completed with error.
+    private(set) var data: Data?
     /// Task completion collected during this interception.
     private(set) var completion: ResourceCompletion?
     /// Trace information propagated with the task. Not available when Tracing is disabled
@@ -35,6 +37,14 @@ internal class TaskInterception {
         self.metrics = metrics
     }
 
+    func register(nextData: Data) {
+        if data != nil {
+            self.data?.append(nextData)
+        } else {
+            self.data = nextData
+        }
+    }
+
     func register(completion: ResourceCompletion) {
         self.completion = completion
     }

Sources/Datadog/URLSessionAutoInstrumentation/Interception/URLSessionInterceptor.swift

@@ -10,6 +10,7 @@ internal protocol URLSessionInterceptorType: class {
     func modify(request: URLRequest, session: URLSession?) -> URLRequest
     func taskCreated(task: URLSessionTask, session: URLSession?)
     func taskMetricsCollected(task: URLSessionTask, metrics: URLSessionTaskMetrics)
+    func taskReceivedData(task: URLSessionTask, data: Data)
     func taskCompleted(task: URLSessionTask, error: Error?)
 }
 
@@ -43,18 +44,20 @@ public class URLSessionInterceptor: URLSessionInterceptorType {
         let handler: URLSessionInterceptionHandler
 
         if configuration.instrumentRUM {
-            handler = URLSessionRUMResourcesHandler(dateProvider: dateProvider)
+            handler = URLSessionRUMResourcesHandler(
+                dateProvider: dateProvider,
+                rumAttributesProvider: configuration.rumAttributesProvider
+            )
         } else {
             handler = URLSessionTracingHandler(appStateListener: appStateListener)
         }
 
-        self.init(configuration: configuration, handler: handler, appStateListener: appStateListener)
+        self.init(configuration: configuration, handler: handler)
     }
 
     init(
         configuration: FeaturesConfiguration.URLSessionAutoInstrumentation,
-        handler: URLSessionInterceptionHandler,
-        appStateListener: AppStateListening
+        handler: URLSessionInterceptionHandler
     ) {
         self.defaultFirstPartyURLsFilter = FirstPartyURLsFilter(hosts: configuration.userDefinedFirstPartyHosts)
         self.internalURLsFilter = InternalURLsFilter(urls: configuration.sdkInternalURLs)
@@ -79,11 +82,11 @@ public class URLSessionInterceptor: URLSessionInterceptorType {
     }
 
     /// An internal queue for synchronising the access to `interceptionByTask`.
-    private let queue = DispatchQueue(label: "com.datadoghq.URLSessionInterceptor", target: .global(qos: .utility))
+    internal let queue = DispatchQueue(label: "com.datadoghq.URLSessionInterceptor", target: .global(qos: .utility))
     /// Maps `URLSessionTask` to its `TaskInterception` object.
     private var interceptionByTask: [URLSessionTask: TaskInterception] = [:]
 
-    // MARK: - Public
+    // MARK: - Interception Flow
 
     /// Intercepts given `URLRequest` before it is sent.
     /// If Tracing feature is enabled and first party hosts are configured in `Datadog.Configuration`, this method will
@@ -127,10 +130,10 @@ public class URLSessionInterceptor: URLSessionInterceptorType {
     }
 
     /// Notifies the `URLSessionTask` metrics collection.
-    /// This method should be called as soon as the task metrics were received by `URLSessionDelegate`.
+    /// This method should be called as soon as the task metrics were received by `URLSessionTaskDelegate`.
     /// - Parameters:
     ///   - task: task receiving metrics.
-    ///   - metrics: metrics object delivered to `URLSessionDelegate`.
+    ///   - metrics: metrics object delivered to `URLSessionTaskDelegate`.
     public func taskMetricsCollected(task: URLSessionTask, metrics: URLSessionTaskMetrics) {
         guard !internalURLsFilter.isInternal(url: task.originalRequest?.url) else {
             return
@@ -151,6 +154,25 @@ public class URLSessionInterceptor: URLSessionInterceptorType {
         }
     }
 
+    /// Notifies the `URLSessionTask` data receiving.
+    /// This method should be called as soon as the next chunk of data is received by `URLSessionDataDelegate`.
+    /// - Parameters:
+    ///   - task: task receiving data.
+    ///   - data: next chunk of data delivered to `URLSessionDataDelegate`.
+    public func taskReceivedData(task: URLSessionTask, data: Data) {
+        guard !internalURLsFilter.isInternal(url: task.originalRequest?.url) else {
+            return
+        }
+
+        queue.async {
+            guard let interception = self.interceptionByTask[task] else {
+                return
+            }
+
+            interception.register(nextData: data)
+        }
+    }
+
     /// Notifies the `URLSessionTask` completion.
     /// This method should be called as soon as the task was completed.
     /// - Parameter task: the task object obtained from `URLSession`.

Sources/Datadog/URLSessionAutoInstrumentation/URLSessionSwizzler.swift

@@ -78,6 +78,9 @@ internal class URLSessionSwizzler {
                         var taskReference: URLSessionDataTask?
                         let newCompletionHandler: CompletionHandler = { data, response, error in
                             if let task = taskReference { // sanity check, should always succeed
+                                if let data = data {
+                                    interceptor.taskReceivedData(task: task, data: data)
+                                }
                                 interceptor.taskCompleted(task: task, error: error)
                             }
                             completionHandler?(data, response, error)
@@ -137,6 +140,9 @@ internal class URLSessionSwizzler {
                         var taskReference: URLSessionDataTask?
                         let newCompletionHandler: CompletionHandler = { data, response, error in
                             if let task = taskReference { // sanity check, should always succeed
+                                if let data = data {
+                                    interceptor.taskReceivedData(task: task, data: data)
+                                }
                                 interceptor.taskCompleted(task: task, error: error)
                             }
                             completionHandler?(data, response, error)

Sources/DatadogExtensions/Alamofire/DatadogAlamofireExtension.swift

@@ -22,6 +22,10 @@ public class DDEventMonitor: EventMonitor {
     public func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) {
         URLSessionInterceptor.shared?.taskCompleted(task: task, error: error)
     }
+
+    public func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive data: Data) {
+        URLSessionInterceptor.shared?.taskReceivedData(task: dataTask, data: data)
+    }
 }
 
 /// An `Alamofire.RequestInterceptor` which instruments `Alamofire.Session` with Datadog RUM and Tracing.

Sources/DatadogObjc/DDURLSessionDelegate+objc.swift

@@ -8,7 +8,7 @@ import Foundation
 import Datadog
 
 @objc
-open class DDNSURLSessionDelegate: NSObject, URLSessionTaskDelegate, __URLSessionDelegateProviding {
+open class DDNSURLSessionDelegate: NSObject, URLSessionTaskDelegate, URLSessionDataDelegate, __URLSessionDelegateProviding {
     var swiftDelegate: DDURLSessionDelegate
     public var ddURLSessionDelegate: DDURLSessionDelegate {
         return swiftDelegate
@@ -31,4 +31,8 @@ open class DDNSURLSessionDelegate: NSObject, URLSessionTaskDelegate, __URLSessio
     open func urlSession(_ session: URLSession, task: URLSessionTask, didFinishCollecting metrics: URLSessionTaskMetrics) {
         swiftDelegate.urlSession(session, task: task, didFinishCollecting: metrics)
     }
+
+    open func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive data: Data) {
+        swiftDelegate.urlSession(session, dataTask: dataTask, didReceive: data)
+    }
 }

Tests/DatadogIntegrationTests/Scenarios/RUM/RUMResourcesScenarioTests.swift

@@ -198,11 +198,25 @@ class RUMResourcesScenarioTests: IntegrationTests, RUMCommonAsserts {
             "Both 3rd party resources should track download phase"
         )
 
-        // Ensure there were no tracing `Spans` send
+        // Assert there were no tracing `Spans` sent
         _ = try tracingServerSession.pullRecordedRequests(timeout: 1) { requests in
             XCTAssertEqual(requests.count, 0, "There should be no tracing `Spans` send")
             return true
         }
+
+        // Assert it adds custom RUM attributes to intercepted RUM Resources:
+        session.resourceEventMatchers.forEach { resourceEvent in
+            XCTAssertNotNil(try? resourceEvent.attribute(forKeyPath: "context.response.body.truncated") as String)
+            XCTAssertNotNil(try? resourceEvent.attribute(forKeyPath: "context.response.headers") as String)
+            XCTAssertNil(try? resourceEvent.attribute(forKeyPath: "context.response.error") as String)
+        }
+
+        // Assert it adds custom RUM attributes to intercepted RUM Resources which finished with error:
+        session.errorEventMatchers.forEach { errorEvent in
+            XCTAssertNil(try? errorEvent.attribute(forKeyPath: "context.response.body.truncated") as String)
+            XCTAssertNil(try? errorEvent.attribute(forKeyPath: "context.response.headers") as String)
+            XCTAssertNotNil(try? errorEvent.attribute(forKeyPath: "context.response.error") as String)
+        }
     }
 
     private func getTraceID(from request: Request) -> String? {