From 720965dfe09400e2ada2431336539cad6e0b6864 Mon Sep 17 00:00:00 2001
From: ArielDemarco <ariel.demarco@embrace.io>
Date: Tue, 24 Dec 2024 12:37:01 -0300
Subject: [PATCH] Ensure thread safe access to `UIViewController` associated
 properties (#153)

---
 .../UX/View/UIViewControllerHandler.swift     | 29 ++++++++++---------
 .../Validators/LengthOfNameValidator.swift    |  7 +++--
 2 files changed, 20 insertions(+), 16 deletions(-)

diff --git a/Sources/EmbraceCore/Capture/UX/View/UIViewControllerHandler.swift b/Sources/EmbraceCore/Capture/UX/View/UIViewControllerHandler.swift
index 53550859..4c26c664 100644
--- a/Sources/EmbraceCore/Capture/UX/View/UIViewControllerHandler.swift
+++ b/Sources/EmbraceCore/Capture/UX/View/UIViewControllerHandler.swift
@@ -85,12 +85,13 @@ class UIViewControllerHandler {
               let otel = dataSource?.otel else {
             return
         }
+        // We generate the id here, outside of the `queue`, to ensure we're doing it while the ViewController is still alive (renerding process).
+        // There could be a race condition and it's possible that the controller was released or is in the process of deallocation,
+        // which could cause a crash (as this feature relies on objc_setAssociatedObject).
+        let id = UUID().uuidString
+        vc.emb_identifier = id
 
         queue.async {
-            // generate id
-            let id = UUID().uuidString
-            vc.emb_identifier = id
-
             // generate parent span
             let className = vc.className
 
@@ -183,17 +184,23 @@ class UIViewControllerHandler {
     }
 
     func onViewDidAppearEnd(_ vc: UIViewController) {
+        if self.dataSource?.instrumentVisibility == true {
+            // Create id only if necessary. This could happen when `instrumentFirstRender` is `false`
+            // in those cases, the `emb_identifier` will be `nil` and we need it to instrument visibility
+            // (and in those cases that's enabled, also instrumenting the rendering process).
+            // The reason why we're doing this outside of the utility `queue` can be found on `onViewDidLoadStart`.
+            if vc.emb_identifier == nil {
+                vc.emb_identifier = UUID().uuidString
+            }
+        }
+
         queue.async {
-            guard let otel = self.dataSource?.otel else {
+            guard let otel = self.dataSource?.otel, let id = vc.emb_identifier else {
                 return
             }
 
             // check if we need to create a visibility span
             if self.dataSource?.instrumentVisibility == true {
-                // create id if necessary
-                let id = vc.emb_identifier ?? UUID().uuidString
-                vc.emb_identifier = id
-
                 let span = self.createSpan(
                     with: otel,
                     vc: vc,
@@ -203,10 +210,6 @@ class UIViewControllerHandler {
                 self.visibilitySpans[id] = span
             }
 
-            guard let id = vc.emb_identifier else {
-                return
-            }
-
             // end view did appear span
             let now = Date()
 
diff --git a/Sources/EmbraceCore/Internal/SpanStorageExporter/Validation/Validators/LengthOfNameValidator.swift b/Sources/EmbraceCore/Internal/SpanStorageExporter/Validation/Validators/LengthOfNameValidator.swift
index ee4d9f97..cd493320 100644
--- a/Sources/EmbraceCore/Internal/SpanStorageExporter/Validation/Validators/LengthOfNameValidator.swift
+++ b/Sources/EmbraceCore/Internal/SpanStorageExporter/Validation/Validators/LengthOfNameValidator.swift
@@ -5,11 +5,12 @@
 import EmbraceOTelInternal
 import EmbraceSemantics
 import OpenTelemetrySdk
+import EmbraceCommonInternal
 
-/// Validates the length of ``SpanData.name``. 
+/// Validates the length of ``SpanData.name``.
 /// This compares the length of the String in characters, not bytes.
 class LengthOfNameValidator: SpanDataValidator {
-
+    private static let allowList: [SpanType] = [.networkRequest, .view, .viewLoad]
     let allowedCharacterCount: ClosedRange<Int>
 
     init(allowedCharacterCount: ClosedRange<Int> = 1...50) {
@@ -24,6 +25,6 @@ class LengthOfNameValidator: SpanDataValidator {
     }
 
     private func shouldValidate(data: SpanData) -> Bool {
-        return data.embType != .networkRequest
+        return !LengthOfNameValidator.allowList.contains(data.embType)
     }
 }