[WIP] Fixed automatic view capture bugs

Sources/EmbraceCommonInternal/Swizzling/EmbraceSwizzler.swift

@@ -0,0 +1,80 @@
+//
+//  Copyright © 2024 Embrace Mobile, Inc. All rights reserved.
+//
+import ObjectiveC.runtime
+
+public class EmbraceSwizzler {
+    public init() {
+    }
+
+    /// Swizzles a specific instance method of a given class.
+    ///
+    /// This method allows you to replace the implementation of an instance method in the specified class (`type`)
+    /// with a custom implementation provided as a block. Only the implementation of the specified class is swizzled,
+    /// and parent class methods with the same selector are not affected.
+    ///
+    /// - Parameters:
+    ///   - type: The class in which the method resides. The method to be swizzled **must belong to this class specifically**.
+    ///   - selector: The selector for the instance method to be swizzled.
+    ///   - implementationType: The expected function signature of the original method (use `@convention(c)`)
+    ///   - blockImplementationType: The expected function signature of the new implementation block (use `@convention(block)`)
+    ///   - block: A closure that accepts the original implementation as input (`implementationType.Type`) and provides the
+    ///   new implementation (`blockImplementationType.Type`)
+    ///
+    /// - Important: This method only operates on methods explicitly declared in the specified class. If the method is inherited from a parent class, **it wont be swizzled**.
+    public func swizzleDeclaredInstanceMethod<T, F>(
+        in type: AnyClass,
+        selector: Selector,
+        implementationType: T.Type,
+        blockImplementationType: F.Type,
+        _ block: @escaping (T) -> F
+    ) throws {
+
+        // Find the method in the specified class
+        var methodToSwizzle: Method?
+        var methodCount: UInt32 = 0
+
+        // We use `class_copyMethodList` and search for the method instead of using `class_getInstanceMethod`
+        // because we don't want to modify the `superclass` implementation.
+        let methods = class_copyMethodList(type, &methodCount)
+        if let methods = methods {
+            for index in 0..<Int(methodCount) {
+                let method = methods[index]
+                if sel_isEqual(method_getName(method), selector) {
+                    methodToSwizzle = method
+                    break
+                }
+            }
+
+        }
+
+        free(methods)
+
+        // If the method is not found, we exit early. This is not a real problem, that's why we don't throw as in `Swizzlable`
+        guard let method = methodToSwizzle else {
+            return
+        }
+
+        // Retrieve the original implementation of the method
+        let originalImplementation = method_getImplementation(method)
+        saveInCache(originalImplementation: originalImplementation, forMethod: method, associatedToClass: type)
+
+        // Create a block implementation by invoking the provided closure, passing the original implementation as input.
+        let originalTypifiedImplementation = unsafeBitCast(originalImplementation, to: implementationType)
+        let newImplementationBlock: F = block(originalTypifiedImplementation)
+        let newImplementation = imp_implementationWithBlock(newImplementationBlock)
+
+        // Do the actual IMP replacement 
+        method_setImplementation(method, newImplementation)
+    }
+
+    private func saveInCache(originalImplementation: IMP, forMethod method: Method, associatedToClass: AnyClass) {
+        #if DEBUG
+        let swizzlerClassName = String(describing: type(of: self))
+        SwizzleCache.shared.addMethodImplementation(originalImplementation,
+                                                    forMethod: method,
+                                                    inClass: associatedToClass,
+                                                    swizzler: swizzlerClassName)
+        #endif
+    }
+}

Sources/EmbraceCore/Capture/UX/View/Protocols/CaptureServices+UIViewController.swift

@@ -64,7 +64,7 @@ extension CaptureServices {
         }
 
         guard let builder = viewCaptureService.otel?.buildSpan(
-            name: name, 
+            name: name,
             type: type,
             attributes: attributes,
             autoTerminationCode: nil

Sources/EmbraceCore/Capture/UX/View/UIViewController+Embrace.swift

@@ -9,22 +9,26 @@ import UIKit
 extension UIViewController {
     private struct AssociatedKeys {
         static var embraceIdentifier: Int = 0
+        static var anotherIdentifier: Int = 1
     }
 
-    var emb_identifier: String? {
+    var emb_instrumentation_state: ViewInstrumentationState? {
         get {
-            if let value = objc_getAssociatedObject(self, &AssociatedKeys.embraceIdentifier) as? NSString {
-                return value as String
+            if let value = objc_getAssociatedObject(
+                self,
+                &AssociatedKeys.anotherIdentifier
+            ) as? ViewInstrumentationState {
+                return value as ViewInstrumentationState
             }
 
             return nil
         }
 
         set {
             objc_setAssociatedObject(self,
-                                     &AssociatedKeys.embraceIdentifier,
+                                     &AssociatedKeys.anotherIdentifier,
                                      newValue,
-                                     .OBJC_ASSOCIATION_RETAIN)
+                                     .OBJC_ASSOCIATION_RETAIN_NONATOMIC)
         }
     }
 

Sources/EmbraceCore/Capture/UX/View/UIViewControllerHandler.swift

@@ -26,6 +26,7 @@ class UIViewControllerHandler {
     @ThreadSafe var parentSpans: [String: Span] = [:]
     @ThreadSafe var viewDidLoadSpans: [String: Span] = [:]
     @ThreadSafe var viewWillAppearSpans: [String: Span] = [:]
+    @ThreadSafe var viewIsAppearingSpans: [String: Span] = [:]
     @ThreadSafe var viewDidAppearSpans: [String: Span] = [:]
     @ThreadSafe var visibilitySpans: [String: Span] = [:]
 
@@ -46,7 +47,7 @@ class UIViewControllerHandler {
     }
 
     func parentSpan(for vc: UIViewController) -> Span? {
-        guard let id = vc.emb_identifier else {
+        guard let id = vc.emb_instrumentation_state?.identifier else {
             return nil
         }
 
@@ -70,14 +71,15 @@ class UIViewControllerHandler {
             self.parentSpans.removeAll()
             self.viewDidLoadSpans.removeAll()
             self.viewWillAppearSpans.removeAll()
+            self.viewIsAppearingSpans.removeAll()
             self.viewDidAppearSpans.removeAll()
             self.visibilitySpans.removeAll()
             self.uiReadySpans.removeAll()
             self.alreadyFinishedUiReadyIds.removeAll()
         }
     }
 
-    func onViewDidLoadStart(_ vc: UIViewController) {
+    func onViewDidLoadStart(_ vc: UIViewController, now: Date = Date()) {
 
         guard dataSource?.state == .active,
               dataSource?.instrumentFirstRender == true,
@@ -89,7 +91,10 @@ class UIViewControllerHandler {
         // 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
+        var state = ViewInstrumentationState()
+        state.viewDidLoadSpanCreated = true
+        state.identifier = id
+        vc.emb_instrumentation_state = state
 
         queue.async {
             // generate parent span
@@ -105,14 +110,16 @@ class UIViewControllerHandler {
             let parentSpan = self.createSpan(
                 with: otel,
                 vc: vc,
-                name: spanName
+                name: spanName,
+                startTime: now
             )
 
             // generate view did load span
             let viewDidLoadSpan = self.createSpan(
                 with: otel,
                 vc: vc,
                 name: SpanSemantics.View.viewDidLoadName,
+                startTime: now,
                 parent: parentSpan
             )
 
@@ -121,21 +128,22 @@ class UIViewControllerHandler {
         }
     }
 
-    func onViewDidLoadEnd(_ vc: UIViewController) {
+    func onViewDidLoadEnd(_ vc: UIViewController, now: Date = Date()) {
         queue.async {
-            guard let id = vc.emb_identifier,
+            guard let id = vc.emb_instrumentation_state?.identifier,
                   let span = self.viewDidLoadSpans.removeValue(forKey: id) else {
                 return
             }
 
-            span.end()
+            span.end(time: now)
         }
     }
 
-    func onViewWillAppearStart(_ vc: UIViewController) {
+    func onViewWillAppearStart(_ vc: UIViewController, now: Date = Date()) {
+        vc.emb_instrumentation_state?.viewWillAppearSpanCreated = true
         queue.async {
             guard let otel = self.dataSource?.otel,
-                  let id = vc.emb_identifier,
+                  let id = vc.emb_instrumentation_state?.identifier,
                   let parentSpan = self.parentSpans[id] else {
                 return
             }
@@ -145,28 +153,63 @@ class UIViewControllerHandler {
                 with: otel,
                 vc: vc,
                 name: SpanSemantics.View.viewWillAppearName,
+                startTime: now,
                 parent: parentSpan
             )
 
             self.viewWillAppearSpans[id] = span
         }
     }
 
-    func onViewWillAppearEnd(_ vc: UIViewController) {
+    func onViewWillAppearEnd(_ vc: UIViewController, now: Date = Date()) {
         queue.async {
-            guard let id = vc.emb_identifier,
+            guard let id = vc.emb_instrumentation_state?.identifier,
                   let span = self.viewWillAppearSpans.removeValue(forKey: id) else {
                 return
             }
 
-            span.end()
+            span.end(time: now)
         }
     }
 
-    func onViewDidAppearStart(_ vc: UIViewController) {
+    func onViewIsAppearingStart(_ vc: UIViewController, now: Date = Date()) {
+        vc.emb_instrumentation_state?.viewIsAppearingSpanCreated = true
         queue.async {
             guard let otel = self.dataSource?.otel,
-                  let id = vc.emb_identifier,
+                  let id = vc.emb_instrumentation_state?.identifier,
+                  let parentSpan = self.parentSpans[id] else {
+                return
+            }
+
+            // generate view is appearing span
+            let span = self.createSpan(
+                with: otel,
+                vc: vc,
+                name: SpanSemantics.View.viewIsAppearingName,
+                startTime: now,
+                parent: parentSpan
+            )
+
+            self.viewIsAppearingSpans[id] = span
+        }
+    }
+
+    func onViewIsAppearingEnd(_ vc: UIViewController, now: Date = Date()) {
+        queue.async {
+            guard let id = vc.emb_instrumentation_state?.identifier,
+                  let span = self.viewIsAppearingSpans.removeValue(forKey: id) else {
+                return
+            }
+
+            span.end(time: now)
+        }
+    }
+
+    func onViewDidAppearStart(_ vc: UIViewController, now: Date = Date()) {
+        vc.emb_instrumentation_state?.viewDidAppearSpanCreated = true
+        queue.async {
+            guard let otel = self.dataSource?.otel,
+                  let id = vc.emb_instrumentation_state?.identifier,
                   let parentSpan = self.parentSpans[id] else {
                 return
             }
@@ -176,26 +219,27 @@ class UIViewControllerHandler {
                 with: otel,
                 vc: vc,
                 name: SpanSemantics.View.viewDidAppearName,
+                startTime: now,
                 parent: parentSpan
             )
 
             self.viewDidAppearSpans[id] = span
         }
     }
 
-    func onViewDidAppearEnd(_ vc: UIViewController) {
+    func onViewDidAppearEnd(_ vc: UIViewController, now: Date = Date()) {
         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
+            if vc.emb_instrumentation_state == nil {
+                vc.emb_instrumentation_state = ViewInstrumentationState(identifier: UUID().uuidString)
             }
         }
 
         queue.async {
-            guard let otel = self.dataSource?.otel, let id = vc.emb_identifier else {
+            guard let otel = self.dataSource?.otel, let id = vc.emb_instrumentation_state?.identifier else {
                 return
             }
 
@@ -205,14 +249,12 @@ class UIViewControllerHandler {
                     with: otel,
                     vc: vc,
                     name: SpanSemantics.View.screenName,
-                    type: .view
+                    type: .view,
+                    startTime: now
                 )
                 self.visibilitySpans[id] = span
             }
 
-            // end view did appear span
-            let now = Date()
-
             if let span = self.viewDidAppearSpans.removeValue(forKey: id) {
                 span.end(time: now)
             }
@@ -232,6 +274,7 @@ class UIViewControllerHandler {
                     with: otel,
                     vc: vc,
                     name: SpanSemantics.View.uiReadyName,
+                    startTime: now,
                     parent: parentSpan
                 )
 
@@ -253,7 +296,7 @@ class UIViewControllerHandler {
 
     func onViewDidDisappear(_ vc: UIViewController) {
         queue.async {
-            guard let id = vc.emb_identifier else {
+            guard let id = vc.emb_instrumentation_state?.identifier else {
                 return
             }
 
@@ -272,7 +315,7 @@ class UIViewControllerHandler {
 
     func onViewBecameInteractive(_ vc: UIViewController) {
         queue.async {
-            guard let id = vc.emb_identifier,
+            guard let id = vc.emb_instrumentation_state?.identifier,
                   let parentSpan = self.parentSpans[id],
                   parentSpan.isTimeToInteractive else {
                 return
@@ -305,6 +348,10 @@ class UIViewControllerHandler {
             viewWillAppearSpan.end(errorCode: .userAbandon, time: time)
         }
 
+        if let viewIsAppearingSpan = self.viewIsAppearingSpans[id] {
+            viewIsAppearingSpan.end(errorCode: .userAbandon, time: time)
+        }
+
         if let viewDidAppearSpan = self.viewDidAppearSpans[id] {
             viewDidAppearSpan.end(errorCode: .userAbandon, time: time)
         }
@@ -325,6 +372,7 @@ class UIViewControllerHandler {
         vc: UIViewController,
         name: String,
         type: SpanType = .viewLoad,
+        startTime: Date,
         parent: Span? = nil
     ) -> Span {
         let builder = otel.buildSpan(
@@ -341,18 +389,20 @@ class UIViewControllerHandler {
             builder.setParent(parent)
         }
 
+        builder.setStartTime(time: startTime)
+
         return builder.startSpan()
     }
 
     private func clear(id: String, vc: UIViewController? = nil) {
         self.parentSpans[id] = nil
         self.viewDidLoadSpans[id] = nil
         self.viewWillAppearSpans[id] = nil
+        self.viewIsAppearingSpans[id] = nil
         self.viewDidAppearSpans[id] = nil
         self.uiReadySpans[id] = nil
         self.alreadyFinishedUiReadyIds.remove(id)
-
-        vc?.emb_identifier = nil
+        vc?.emb_instrumentation_state = nil
     }
 }