@@ -497,8 +497,10 @@ public class URLSessionInstrumentation {
497497 // URLSessionTask methods
498498 private func urlSession( _ session: URLSession , dataTask: URLSessionDataTask , didReceive data: Data ) {
499499 guard configuration. shouldRecordPayload ? ( session) ?? false else { return }
500+ guard let taskId = objc_getAssociatedObject ( dataTask, & idKey) as? String else {
501+ return
502+ }
500503 let dataCopy = data
501- let taskId = idKeyForTask ( dataTask)
502504 queue. sync {
503505 if ( requestMap [ taskId] ? . request) != nil {
504506 createRequestState ( for: taskId)
@@ -512,7 +514,9 @@ public class URLSessionInstrumentation {
512514
513515 private func urlSession( _ session: URLSession , dataTask: URLSessionDataTask , didReceive response: URLResponse , completionHandler: @escaping ( URLSession . ResponseDisposition ) -> Void ) {
514516 guard configuration. shouldRecordPayload ? ( session) ?? false else { return }
515- let taskId = idKeyForTask ( dataTask)
517+ guard let taskId = objc_getAssociatedObject ( dataTask, & idKey) as? String else {
518+ return
519+ }
516520 queue. sync {
517521 if ( requestMap [ taskId] ? . request) != nil {
518522 createRequestState ( for: taskId)
@@ -526,7 +530,9 @@ public class URLSessionInstrumentation {
526530 }
527531
528532 private func urlSession( _ session: URLSession , task: URLSessionTask , didCompleteWithError error: Error ? ) {
529- let taskId = idKeyForTask ( task)
533+ guard let taskId = objc_getAssociatedObject ( task, & idKey) as? String else {
534+ return
535+ }
530536 var requestState : NetworkRequestState ?
531537 queue. sync {
532538 requestState = requestMap [ taskId]
@@ -543,13 +549,16 @@ public class URLSessionInstrumentation {
543549 }
544550
545551 private func urlSession( _ session: URLSession , dataTask: URLSessionDataTask , didBecome downloadTask: URLSessionDownloadTask ) {
546- let id = idKeyForTask ( dataTask)
547- setIdKey ( value: id, for: downloadTask)
552+ guard let taskId = objc_getAssociatedObject ( dataTask, & idKey) as? String else {
553+ return
554+ }
555+ self . setIdKey ( value: taskId, for: downloadTask)
548556 }
549557
550558 private func urlSession( _ session: URLSession , task: URLSessionTask , didFinishCollecting metrics: URLSessionTaskMetrics ) {
551- let taskId = idKeyForTask ( task)
552-
559+ guard let taskId = objc_getAssociatedObject ( task, & idKey) as? String else {
560+ return
561+ }
553562 var requestState : NetworkRequestState ?
554563 queue. sync {
555564 requestState = requestMap [ taskId]
0 commit comments