From 61326afdae44384cab3e56015c77d989b7be5b0a Mon Sep 17 00:00:00 2001
From: Jonah Williams <jonahwilliams@google.com>
Date: Tue, 2 Jul 2024 08:34:04 -0700
Subject: [PATCH] [flutter_driver] add allocator mtl to memory event allowlist.
 (#151153)

Flutter driver side changes for https://github.com/flutter/engine/pull/53618
---
 .../lib/src/driver/memory_summarizer.dart     |  2 +-
 .../test/src/gpu_memory_summarizer_test.dart  | 21 +++++++++++++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/packages/flutter_driver/lib/src/driver/memory_summarizer.dart b/packages/flutter_driver/lib/src/driver/memory_summarizer.dart
index 20fb9e40baa..58a3c8ca5c9 100644
--- a/packages/flutter_driver/lib/src/driver/memory_summarizer.dart
+++ b/packages/flutter_driver/lib/src/driver/memory_summarizer.dart
@@ -21,7 +21,7 @@ class GPUMemorySumarizer {
   }
 
   /// Whether or not this event is a GPU allocation event.
-  static const Set<String> kMemoryEvents = <String>{'AllocatorVK'};
+  static const Set<String> kMemoryEvents = <String>{'AllocatorVK', 'AllocatorMTL'};
 
   final List<double> _memoryMB = <double>[];
 
diff --git a/packages/flutter_driver/test/src/gpu_memory_summarizer_test.dart b/packages/flutter_driver/test/src/gpu_memory_summarizer_test.dart
index 0dc4220bdee..56d67cbbe4c 100644
--- a/packages/flutter_driver/test/src/gpu_memory_summarizer_test.dart
+++ b/packages/flutter_driver/test/src/gpu_memory_summarizer_test.dart
@@ -15,6 +15,14 @@ TimelineEvent newGPUTraceEvent(double ms) => TimelineEvent(<String, dynamic>{
   },
 });
 
+TimelineEvent newMetalGPUTraceEvent(double ms) => TimelineEvent(<String, dynamic>{
+  'name': 'AllocatorMTL',
+  'ph': 'b',
+  'args': <String, String>{
+    'MemoryBudgetUsageMB': ms.toString()
+  },
+});
+
 void main() {
   test('Can process GPU memory usage times.', () {
     final GPUMemorySumarizer summarizer = GPUMemorySumarizer(<TimelineEvent>[
@@ -28,4 +36,17 @@ void main() {
     expect(summarizer.computePercentileMemoryUsage(50.0), closeTo(1024, 0.1));
     expect(summarizer.computeWorstMemoryUsage(), 2048);
   });
+
+  test('Can process Metal GPU memory usage times.', () {
+    final GPUMemorySumarizer summarizer = GPUMemorySumarizer(<TimelineEvent>[
+      newMetalGPUTraceEvent(1024),
+      newMetalGPUTraceEvent(1024),
+      newMetalGPUTraceEvent(512),
+      newMetalGPUTraceEvent(2048),
+    ]);
+
+    expect(summarizer.computeAverageMemoryUsage(), closeTo(1152, 0.1));
+    expect(summarizer.computePercentileMemoryUsage(50.0), closeTo(1024, 0.1));
+    expect(summarizer.computeWorstMemoryUsage(), 2048);
+  });
 }