From 6dbeeb34e19c3c1f3a177178336cd450236813a6 Mon Sep 17 00:00:00 2001
From: Maxime Epain <maxime.epain@datadoghq.com>
Date: Wed, 14 Aug 2024 14:44:06 +0200
Subject: [PATCH] RUM-5555 Collect CPU usage

---
 .../Benchmarks/Sources/Benchmarks.swift       | 13 ++++
 .../Benchmarks/Sources/Metrics.swift          | 59 +++++++++++++++++--
 2 files changed, 67 insertions(+), 5 deletions(-)

diff --git a/BenchmarkTests/Benchmarks/Sources/Benchmarks.swift b/BenchmarkTests/Benchmarks/Sources/Benchmarks.swift
index e1338cb737..991dc0c468 100644
--- a/BenchmarkTests/Benchmarks/Sources/Benchmarks.swift
+++ b/BenchmarkTests/Benchmarks/Sources/Benchmarks.swift
@@ -128,6 +128,19 @@ public enum Benchmarks {
             }
         }
 
+        _ = meter.createDoubleObservableGauge(name: "ios.benchmark.cpu") { metric in
+            do {
+                let usage = try CPU.usage()
+                metric.observe(value: usage, labels: labels)
+            } catch {
+                logger.logRecordBuilder()
+                    .setSeverity(.error)
+                    .setAttributes(labels.mapValues { .string($0) })
+                    .setBody("Failed to read CPU Metric: \(error)")
+                    .emit()
+            }
+        }
+
         OpenTelemetry.registerMeterProvider(meterProvider: meterProvider)
     }
 }
diff --git a/BenchmarkTests/Benchmarks/Sources/Metrics.swift b/BenchmarkTests/Benchmarks/Sources/Metrics.swift
index f8ca08d6f3..1bef64a408 100644
--- a/BenchmarkTests/Benchmarks/Sources/Metrics.swift
+++ b/BenchmarkTests/Benchmarks/Sources/Metrics.swift
@@ -12,7 +12,9 @@ let TASK_VM_INFO_COUNT = mach_msg_type_number_t(MemoryLayout<task_vm_info_data_t
 let TASK_VM_INFO_REV1_COUNT = mach_msg_type_number_t(MemoryLayout.offset(of: \task_vm_info_data_t.min_address)! / MemoryLayout<integer_t>.size)
 
 enum MachError: Error {
-    case kernelError(kern_return_t)
+    case task_info(return: kern_return_t)
+    case task_threads(return: kern_return_t)
+    case thread_info(return: kern_return_t)
 }
 
 public enum Memory {
@@ -26,15 +28,62 @@ public enum Memory {
     static func footprint() throws -> Double {
         var info = task_vm_info_data_t()
         var count = TASK_VM_INFO_COUNT
-        let kr = withUnsafeMutablePointer(to: &info) { infoPtr in
-            infoPtr.withMemoryRebound(to: integer_t.self, capacity: Int(count)) { intPtr in
-                task_info(mach_task_self_, task_flavor_t(TASK_VM_INFO), intPtr, &count)
+        let kr = withUnsafeMutablePointer(to: &info) {
+            $0.withMemoryRebound(to: integer_t.self, capacity: Int(count)) {
+                task_info(mach_task_self_, task_flavor_t(TASK_VM_INFO), $0, &count)
             }
         }
+
         guard kr == KERN_SUCCESS, count >= TASK_VM_INFO_REV1_COUNT else {
-            throw MachError.kernelError(kr)
+            throw MachError.task_info(return: kr)
         }
 
         return Double(info.phys_footprint)
     }
 }
+
+public enum CPU {
+    /// Collect single sample of current cpu usage.
+    ///
+    /// The computation is based on https://gist.github.com/hisui/10004131#file-cpu-usage-cpp
+    /// It reads the `cpu_usage` from all thread to compute the application usage percentage.
+    ///
+    /// - Returns: The cpu usage of all threads.
+    static func usage() throws -> Double {
+        var threads_list: thread_act_array_t?
+        var threads_count = mach_msg_type_number_t()
+        let kr = withUnsafeMutablePointer(to: &threads_list) {
+            $0.withMemoryRebound(to: thread_act_array_t?.self, capacity: 1) {
+                task_threads(mach_task_self_, $0, &threads_count)
+            }
+        }
+
+        guard kr == KERN_SUCCESS, let threads_list = threads_list else {
+            throw MachError.task_threads(return: kr)
+        }
+
+        defer {
+            vm_deallocate(mach_task_self_, vm_address_t(bitPattern: threads_list), vm_size_t(Int(threads_count) * MemoryLayout<thread_t>.stride))
+        }
+
+        return try (0..<threads_count).reduce(0) { result, index in
+            var basic_info = thread_basic_info()
+            var basic_info_count = mach_msg_type_number_t(THREAD_INFO_MAX)
+            let kr = withUnsafeMutablePointer(to: &basic_info) {
+                $0.withMemoryRebound(to: integer_t.self, capacity: 1) {
+                    thread_info(threads_list[Int(index)], thread_flavor_t(THREAD_BASIC_INFO), $0, &basic_info_count)
+                }
+            }
+
+            guard kr == KERN_SUCCESS else {
+                throw MachError.thread_info(return: kr)
+            }
+
+            guard basic_info.flags != TH_FLAGS_IDLE else {
+                return result
+            }
+
+            return result + Double(basic_info.cpu_usage) / Double(TH_USAGE_SCALE)
+        }
+    }
+}