Use Concurrency for multithread processing

ETTrace/ETTraceRunner/RunnerHelper.swift

@@ -13,6 +13,8 @@ import Swifter
 import JSONWrapper
 import ETModels
 
+typealias FlamegraphTuple = (String, Thread, Flamegraph)
+
 class RunnerHelper {
     let dsyms: String?
     let launch: Bool
@@ -90,15 +92,25 @@ class RunnerHelper {
         symbolicator = Symbolicator(isSimulator: isSimulator, dSymsDir: dsyms, osVersion: osVersion, arch: arch, verbose: verbose)
         let outputUrl = URL(fileURLWithPath: FileManager.default.currentDirectoryPath)
 
-        var allThreads:[Flamegraph] = []
+        let flamegraphs = await withTaskGroup(of: FlamegraphTuple.self, returning: [FlamegraphTuple].self) { taskGroup in
+            for (threadId, thread) in responseData.threads {
+                taskGroup.addTask {
+                    (threadId, thread, await self.createFlamegraphForThread(thread, responseData))
+                }
+            }
+
+            var tuples = [FlamegraphTuple]()
+            for await result in taskGroup {
+                tuples.append(result)
+            }
+            return tuples
+        }
+
         var mainThreadFlamegraph: Flamegraph!
         var mainThreadData: Data!
-        for (threadId, thread) in responseData.threads {
-            let flamegraph = createFlamegraphForThread(thread, responseData)
-            allThreads.append(flamegraph)
-
+
+        for (threadId, thread, flamegraph) in flamegraphs {
             let outJsonData = JSONWrapper.toData(flamegraph)!
-
             if thread.name == "Main Thread" {
                 mainThreadFlamegraph = flamegraph
                 mainThreadData = outJsonData
@@ -121,9 +133,9 @@ class RunnerHelper {
         print("Results saved to \(outputUrl)")
     }
 
-    private func createFlamegraphForThread(_ thread: Thread, _ responseData: ResponseModel) -> Flamegraph {
+    private func createFlamegraphForThread(_ thread: Thread, _ responseData: ResponseModel) async -> Flamegraph {
         let stacks = thread.stacks
-        let syms = symbolicator.symbolicate(stacks, responseData.libraryInfo.loadedLibraries)
+        let syms = await symbolicator.symbolicate(stacks, responseData.libraryInfo.loadedLibraries)
         let flamegraphNodes = FlamegraphGenerator.generateFlamegraphs(stacks: stacks, syms: syms, writeFolded: verbose)
         let threadNode = ThreadNode(nodes: flamegraphNodes, threadName: thread.name)
 

ETTrace/ETTraceRunner/Utils/Sequence+AsyncMap.swift

@@ -0,0 +1,22 @@
+//
+//  Sequence+AsyncMap.swift
+//  
+//
+//  Created by Itay Brenner on 26/9/23.
+//
+
+import Foundation
+
+extension Sequence {
+    func asyncMap<T>(
+        _ transform: (Element) async throws -> T
+    ) async rethrows -> [T] {
+        var values = [T]()
+
+        for element in self {
+            try await values.append(transform(element))
+        }
+
+        return values
+    }
+}

ETTrace/ETTraceRunner/Utils/Symbolicator.swift

@@ -12,14 +12,27 @@ struct Address {
     let lib: LoadedLibrary?
 }
 
-class Symbolicator {
+// Actor to help concurrency access
+fileprivate actor SymbolicatorHepler {
+    var libToAddrToSym: [String: [UInt64: String]] = [:]
     var formatSymbolCache: [String: String] = [:]
 
+    func setValue(_ lib: String, addrToSym: [UInt64: String]) {
+        libToAddrToSym[lib] = addrToSym
+    }
+
+    func setCacheValue(_ sym: String, _ result: String) {
+        formatSymbolCache[sym] = result
+    }
+}
+
+class Symbolicator {
     let isSimulator: Bool
     let dSymsDir: String?
     let osVersion: String
     let arch: String
     let verbose: Bool
+    private var helper: SymbolicatorHepler = SymbolicatorHepler()
 
     init(isSimulator: Bool, dSymsDir: String?, osVersion: String, arch: String, verbose: Bool) {
         self.isSimulator = isSimulator
@@ -29,7 +42,7 @@ class Symbolicator {
         self.verbose = verbose
     }
 
-    func symbolicate(_ stacks: [Stack], _ loadedLibs: [LoadedLibrary]) -> [[(String?, String, UInt64?)]] {
+    func symbolicate(_ stacks: [Stack], _ loadedLibs: [LoadedLibrary]) async -> [[(String?, String, UInt64?)]] {
         var libToAddrs: [LoadedLibrary: Set<UInt64>] = [:]
         let stacks = stacksFromResults(stacks, loadedLibs)
         stacks.flatMap { $0 }.forEach { addr in
@@ -38,39 +51,34 @@ class Symbolicator {
             }
         }
 
-        let stateLock = NSLock()
         var libToCleanedPath = [String: (String, String)]()
-        var libToAddrToSym: [String: [UInt64: String]] = [:]
-        let queue = DispatchQueue(label: "com.emerge.symbolication", qos: .userInitiated, attributes: .concurrent)
-        let group = DispatchGroup()
-        for (lib, addrs) in libToAddrs {
-            let cleanedPath = cleanedUpPath(lib.path)
-            libToCleanedPath[lib.path] = (cleanedPath, URL(string: cleanedPath)?.lastPathComponent ?? "")
-            group.enter()
-            queue.async {
+
+        await withTaskGroup(of: Void
+            .self) { [unowned self] taskGroup in
+            for (lib, addrs) in libToAddrs {
+                let cleanedPath = cleanedUpPath(lib.path)
+                libToCleanedPath[lib.path] = (cleanedPath, URL(string: cleanedPath)?.lastPathComponent ?? "")
+
                 if let dSym = self.dsymForLib(lib) {
-                    let addrToSym = Self.addrToSymForBinary(dSym, addrs)
-                    stateLock.lock()
-                    libToAddrToSym[lib.path] = addrToSym
-                    stateLock.unlock()
+                    taskGroup.addTask {
+                        await self.recordAddresToSym(dSym, addrs, lib.path)
+                    }
                 }
-                group.leave()
             }
         }
-        group.wait()
 
         var noLibCount = 0
         var noSymMap: [String: UInt64] = [:]
-        let result: [[(String?, String, UInt64?)]] = stacks.map { stack in
-            stack.map { addr in
+        let result: [[(String?, String, UInt64?)]] = await stacks.asyncMap { stack in
+            await stack.asyncMap { addr in
                 if let lib = addr.lib {
                     let (libPath, lastPathComponent) = libToCleanedPath[lib.path]!
-                    guard let addrToSym = libToAddrToSym[lib.path],
+                    guard let addrToSym = await helper.libToAddrToSym[lib.path],
                           let sym = addrToSym[addr.addr] else {
                         noSymMap[libPath, default: 0] += 1
                       return (libPath, lastPathComponent, addr.addr)
                     }
-                    return (libPath, formatSymbol(sym), nil)
+                    return (libPath, await formatSymbol(sym), nil)
                 } else {
                     noLibCount += 1
                     return ("<unknown>", "<unknown>", nil)
@@ -89,6 +97,11 @@ class Symbolicator {
         return result
     }
 
+    private func recordAddresToSym(_ dSym: String, _ addrs: Set<UInt64>, _ lib: String) async {
+        let addrToSym = self.addrToSymForBinary(dSym, addrs)
+        await helper.setValue(lib, addrToSym: addrToSym)
+    }
+
     private func cleanedUpPath(_ path: String) -> String {
         if path.contains(".app/") && !path.contains("/Xcode.app/") {
             return path.split(separator: "/").drop(while: { $0.hasSuffix(".app") }).joined(separator: "/")
@@ -137,7 +150,7 @@ class Symbolicator {
         return addrs
     }
 
-    private static func addrToSymForBinary(_ binary: String, _ addrs: Set<UInt64>) -> [UInt64: String] {
+    private func addrToSymForBinary(_ binary: String, _ addrs: Set<UInt64>) -> [UInt64: String] {
         let addrsArray = Array(addrs)
         let addrsFile = NSURL(fileURLWithPath: NSTemporaryDirectory()).appendingPathComponent(UUID().uuidString)!.path
 
@@ -170,8 +183,8 @@ class Symbolicator {
         return result
     }
 
-    private func formatSymbol(_ sym: String) -> String {
-        if let cachedResult = formatSymbolCache[sym] {
+    private func formatSymbol(_ sym: String) async -> String {
+        if let cachedResult = await helper.formatSymbolCache[sym] {
             return cachedResult
         }
         let result = sym.replacingOccurrences(of: ":\\d+\\)", with: ")", options: .regularExpression) // static AppDelegate.$main() (in emergeTest) (AppDelegate.swift:10)
@@ -182,7 +195,7 @@ class Symbolicator {
             .replacingOccurrences(of: "^__\\d+\\+", with: "", options: .regularExpression)
             .replacingOccurrences(of: "^__\\d+\\-", with: "", options: .regularExpression)
             .trimmingCharacters(in: .whitespacesAndNewlines)
-        formatSymbolCache[sym] = result
+        await helper.setCacheValue(sym, result)
         return result
     }