Skip to content

Commit 742132f

Browse files
author
Ignacio Bonafonte
authored
Merge pull request #231 from nachoBonafonte/fix-URLSessionInstrumentation-ios12
2 parents 4766969 + ff3e6b6 commit 742132f

File tree

4 files changed

+96
-32
lines changed

4 files changed

+96
-32
lines changed

‎Sources/Exporters/DatadogExporter/Utils/EncodableValue.swift‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ internal struct JSONStringEncodableValue: Encodable {
6262
} else {
6363
let jsonData: Data
6464

65-
if #available(iOS 13.0, *) {
65+
if #available(OSX 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *) {
6666
jsonData = try jsonEncoder.encode(encodable)
6767
} else {
6868
// Prior to `iOS13.0` the `JSONEncoder` is unable to encode primitive values - it expects them to be

‎Sources/Instrumentation/URLSession/URLSessionInstrumentation.swift‎

Lines changed: 92 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,14 @@ import OpenTelemetrySdk
1010
struct NetworkRequestState {
1111
var request: URLRequest?
1212
var dataProcessed: Data?
13+
14+
mutating func setRequest(_ request: URLRequest) {
15+
self.request = request
16+
}
17+
18+
mutating func setData(_ data: URLRequest) {
19+
self.request = data
20+
}
1321
}
1422

1523
private var idKey: Void?
@@ -68,10 +76,13 @@ public class URLSessionInstrumentation {
6876
}
6977
}
7078
}
71-
injectIntoNSURLSessionCreateTaskMethods()
79+
if #available(OSX 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *) {
80+
injectIntoNSURLSessionCreateTaskMethods()
81+
}
7282
injectIntoNSURLSessionCreateTaskWithParameterMethods()
7383
injectIntoNSURLSessionAsyncDataAndDownloadTaskMethods()
7484
injectIntoNSURLSessionAsyncUploadTaskMethods()
85+
injectIntoNSURLSessionTaskResume()
7586
}
7687

7788
private func injectIntoDelegateClass(cls: AnyClass) {
@@ -207,23 +218,26 @@ public class URLSessionInstrumentation {
207218
var task: URLSessionTask!
208219

209220
var completionBlock = completion
210-
if objc_getAssociatedObject(argument, &idKey) == nil {
211-
let completionWrapper: (Any?, URLResponse?, Error?) -> Void = { object, response, error in
212-
if error != nil {
213-
let status = (response as? HTTPURLResponse)?.statusCode ?? 0
214-
URLSessionLogger.logError(error!, dataOrFile: object, statusCode: status, instrumentation: self, sessionTaskId: sessionTaskId)
215-
} else {
216-
if let response = response {
217-
URLSessionLogger.logResponse(response, dataOrFile: object, instrumentation: self, sessionTaskId: sessionTaskId)
221+
222+
if completionBlock != nil {
223+
if objc_getAssociatedObject(argument, &idKey) == nil {
224+
let completionWrapper: (Any?, URLResponse?, Error?) -> Void = { object, response, error in
225+
if error != nil {
226+
let status = (response as? HTTPURLResponse)?.statusCode ?? 0
227+
URLSessionLogger.logError(error!, dataOrFile: object, statusCode: status, instrumentation: self, sessionTaskId: sessionTaskId)
228+
} else {
229+
if let response = response {
230+
URLSessionLogger.logResponse(response, dataOrFile: object, instrumentation: self, sessionTaskId: sessionTaskId)
231+
}
232+
}
233+
if let completion = completion {
234+
completion(object, response, error)
235+
} else {
236+
(session.delegate as? URLSessionTaskDelegate)?.urlSession?(session, task: task, didCompleteWithError: error)
218237
}
219238
}
220-
if let completion = completion {
221-
completion(object, response, error)
222-
} else {
223-
(session.delegate as? URLSessionTaskDelegate)?.urlSession?(session, task: task, didCompleteWithError: error)
224-
}
239+
completionBlock = completionWrapper
225240
}
226-
completionBlock = completionWrapper
227241
}
228242

229243
if let request = argument as? URLRequest, objc_getAssociatedObject(argument, &idKey) == nil {
@@ -295,6 +309,45 @@ public class URLSessionInstrumentation {
295309
}
296310
}
297311

312+
private func injectIntoNSURLSessionTaskResume() {
313+
var methodsToSwizzle = [Method]()
314+
315+
if let method = class_getInstanceMethod(URLSessionTask.self, #selector(URLSessionTask.resume)) {
316+
methodsToSwizzle.append(method)
317+
}
318+
319+
if let cfURLSession = NSClassFromString("__NSCFURLSessionTask"),
320+
let method = class_getInstanceMethod(cfURLSession, NSSelectorFromString("resume"))
321+
{
322+
methodsToSwizzle.append(method)
323+
}
324+
325+
if NSClassFromString("AFURLSessionManager") != nil {
326+
let classes = InstrumentationUtils.objc_getClassList()
327+
classes.forEach {
328+
if let method = class_getInstanceMethod($0, NSSelectorFromString("af_resume")) {
329+
methodsToSwizzle.append(method)
330+
}
331+
}
332+
}
333+
334+
methodsToSwizzle.forEach {
335+
let theMethod = $0
336+
337+
var originalIMP: IMP?
338+
let block: @convention(block) (URLSessionTask) -> Void = { anyTask in
339+
self.urlSessionTaskWillResume(anyTask)
340+
let key = String(theMethod.hashValue)
341+
objc_setAssociatedObject(anyTask, key, true, .OBJC_ASSOCIATION_RETAIN_NONATOMIC)
342+
let castedIMP = unsafeBitCast(originalIMP, to: (@convention(c) (Any) -> Void).self)
343+
castedIMP(anyTask)
344+
objc_setAssociatedObject(anyTask, key, nil, .OBJC_ASSOCIATION_RETAIN_NONATOMIC)
345+
}
346+
let swizzledIMP = imp_implementationWithBlock(unsafeBitCast(block, to: AnyObject.self))
347+
originalIMP = method_setImplementation(theMethod, swizzledIMP)
348+
}
349+
}
350+
298351
// Delegate methods
299352
private func injectTaskDidReceiveDataIntoDelegateClass(cls: AnyClass) {
300353
let selector = #selector(URLSessionDataDelegate.urlSession(_:dataTask:didReceive:))
@@ -429,8 +482,11 @@ public class URLSessionInstrumentation {
429482
queue.async {
430483
let taskId = self.idKeyForTask(dataTask)
431484
if (self.requestMap[taskId]?.request) != nil {
432-
var requestState = self.requestState(for: taskId)
433-
requestState.dataProcessed?.append(dataCopy)
485+
self.createRequestState(for: taskId)
486+
if self.requestMap[taskId]?.dataProcessed == nil {
487+
self.requestMap[taskId]?.dataProcessed = Data()
488+
}
489+
self.requestMap[taskId]?.dataProcessed?.append(dataCopy)
434490
}
435491
}
436492
}
@@ -440,11 +496,11 @@ public class URLSessionInstrumentation {
440496
queue.async {
441497
let taskId = self.idKeyForTask(dataTask)
442498
if (self.requestMap[taskId]?.request) != nil {
443-
var requestState = self.requestState(for: taskId)
499+
self.createRequestState(for: taskId)
444500
if response.expectedContentLength < 0 {
445-
requestState.dataProcessed = Data()
501+
self.requestMap[taskId]?.dataProcessed = Data()
446502
} else {
447-
requestState.dataProcessed = Data(capacity: Int(response.expectedContentLength))
503+
self.requestMap[taskId]?.dataProcessed = Data(capacity: Int(response.expectedContentLength))
448504
}
449505
}
450506
}
@@ -453,11 +509,8 @@ public class URLSessionInstrumentation {
453509
private func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) {
454510
let taskId = self.idKeyForTask(task)
455511

456-
var requestState: NetworkRequestState?
457-
if (self.requestMap[taskId]?.request) != nil {
458-
requestState = self.requestState(for: taskId)
459-
}
460-
512+
let requestState = self.requestMap[taskId]
513+
461514
if let error = error {
462515
let status = (task.response as? HTTPURLResponse)?.statusCode ?? 0
463516
URLSessionLogger.logError(error, dataOrFile: requestState?.dataProcessed, statusCode: status, instrumentation: self, sessionTaskId: taskId)
@@ -483,6 +536,18 @@ public class URLSessionInstrumentation {
483536
}
484537
}
485538

539+
private func urlSessionTaskWillResume(_ session: URLSessionTask) {
540+
let taskId = self.idKeyForTask(session)
541+
if let request = session.currentRequest {
542+
var state = requestMap[taskId]
543+
if state == nil {
544+
state = NetworkRequestState()
545+
requestMap[taskId] = state
546+
}
547+
requestMap[taskId]?.setRequest(request)
548+
}
549+
}
550+
486551
// Helpers
487552
private func idKeyForTask(_ task: URLSessionTask) -> String {
488553
var id = objc_getAssociatedObject(task, &idKey) as? String
@@ -497,12 +562,11 @@ public class URLSessionInstrumentation {
497562
objc_setAssociatedObject(task, &idKey, value, .OBJC_ASSOCIATION_RETAIN_NONATOMIC)
498563
}
499564

500-
private func requestState(for id: String) -> NetworkRequestState {
565+
private func createRequestState(for id: String) {
501566
var state = requestMap[id]
502-
if state == nil {
567+
if requestMap[id] == nil {
503568
state = NetworkRequestState()
504569
requestMap[id] = state
505570
}
506-
return state!
507571
}
508572
}

‎Sources/Instrumentation/URLSession/URLSessionLogger.swift‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,7 @@ class URLSessionLogger {
7979
}
8080

8181
var returnRequest: URLRequest?
82-
if shouldInjectHeaders {
82+
if shouldInjectHeaders && (instrumentation.configuration.shouldInjectTracingHeaders?(request) ?? true) {
8383
returnRequest = instrumentedRequest(for: request, span: span, instrumentation: instrumentation)
8484
}
8585

@@ -91,7 +91,7 @@ class URLSessionLogger {
9191

9292
instrumentation.configuration.createdRequest?(returnRequest ?? request, span)
9393

94-
return returnRequest ?? request
94+
return returnRequest
9595
}
9696

9797
/// This methods ends a Span when a response arrives

‎Tests/ExportersTests/DatadogExporter/Utils/JSONEncoderTests.swift‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ class JSONEncoderTests: XCTestCase {
2222
EncodingContainer(URL(string: "https://example.com/foo")!)
2323
)
2424

25-
if #available(iOS 13.0, OSX 10.15, *) {
25+
if #available(OSX 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *) {
2626
XCTAssertEqual(encodedURL.utf8String, #"{"value":"https://example.com/foo"}"#)
2727
} else {
2828
XCTAssertEqual(encodedURL.utf8String, #"{"value":"https:\/\/example.com\/foo"}"#)

0 commit comments

Comments
 (0)