ui: tp: Add DrawFrames and SurfaceFlinger to AndroidFrames extension

src/trace_processor/perfetto_sql/stdlib/android/input.sql

diff --git a/src/trace_processor/perfetto_sql/stdlib/android/input.sql b/src/trace_processor/perfetto_sql/stdlib/android/input.sql
index 09d2d67..8785ee2 100644
--- a/src/trace_processor/perfetto_sql/stdlib/android/input.sql
+++ b/src/trace_processor/perfetto_sql/stdlib/android/input.sql
@@ -15,6 +15,8 @@
 
 INCLUDE PERFETTO MODULE android.frames.timeline;
 
+INCLUDE PERFETTO MODULE android.surfaceflinger;
+
 INCLUDE PERFETTO MODULE intervals.intersect;
 
 INCLUDE PERFETTO MODULE slices.with_context;
@@ -607,13 +609,65 @@
 WHERE
   $slice_id IN (s_read.id, s_disp.id, s_recv.id, s_cons.id);
 
-CREATE PERFETTO VIEW _android_input_frames AS
+CREATE PERFETTO TABLE _input_sf_composite AS
 SELECT
-  chor.upid,
-  chor.frame_id,
+  s.id,
+  s.track_id,
+  s.ts,
+  s.dur,
+  cast_int!(STR_SPLIT(s.name, ' ', 1)) AS vsync_id
+FROM slice AS s
+JOIN thread_track AS t
+  ON s.track_id = t.id
+JOIN thread USING (utid)
+JOIN process USING (upid)
+WHERE
+  process.name = '/system/bin/surfaceflinger'
+  AND thread.is_main_thread
+  AND s.name GLOB 'composite *';
+
+CREATE PERFETTO INDEX _input_sf_composite_idx ON _input_sf_composite(vsync_id);
+
+CREATE PERFETTO TABLE _input_sf_resolved AS
+SELECT m.app_upid, m.app_vsync, s.id, s.track_id, s.ts, s.dur
+FROM android_app_to_sf_frame_timeline_match AS m
+JOIN _input_sf_composite AS s ON s.vsync_id = m.sf_vsync;
+
+CREATE PERFETTO INDEX _input_sf_resolved_idx ON _input_sf_resolved(id);
+
+CREATE PERFETTO VIEW _android_input_frames AS
+WITH
+  all_frame_keys AS (
+    SELECT frame_id, upid FROM android_frames_choreographer_do_frame
+    UNION
+    SELECT frame_id, upid FROM android_frames_draw_frame
+    UNION
+    SELECT app_vsync AS frame_id, app_upid AS upid FROM _input_sf_resolved
+  )
+SELECT
+  k.upid,
+  k.frame_id,
   chor.id AS id_do_frame,
   s_chor.track_id AS track_do_frame,
   chor.ts AS ts_do_frame,
-  s_chor.dur AS dur_do_frame
-FROM android_frames_choreographer_do_frame AS chor
-JOIN slice AS s_chor ON s_chor.id = chor.id;
+  s_chor.dur AS dur_do_frame,
+  d.id AS id_draw_frames,
+  s_draw.track_id AS track_draw_frames,
+  s_draw.ts AS ts_draw_frames,
+  s_draw.dur AS dur_draw_frames,
+  sf.id AS id_sf,
+  sf.track_id AS track_sf,
+  sf.ts AS ts_sf,
+  sf.dur AS dur_sf
+FROM all_frame_keys AS k
+LEFT JOIN android_frames_choreographer_do_frame AS chor
+  ON chor.frame_id = k.frame_id
+  AND chor.upid = k.upid
+LEFT JOIN slice AS s_chor ON s_chor.id = chor.id
+LEFT JOIN android_frames_draw_frame AS d
+  ON d.frame_id = k.frame_id
+  AND d.upid = k.upid
+LEFT JOIN slice AS s_draw ON s_draw.id = d.id
+LEFT JOIN _input_sf_resolved AS sf
+  ON sf.app_vsync = k.frame_id
+  AND sf.app_upid = k.upid;

ui/src/plugins/com.android.AndroidInputLifecycle/extensions/android_frames_extension.ts

diff --git a/ui/src/plugins/com.android.AndroidInputLifecycle/extensions/android_frames_extension.ts b/ui/src/plugins/com.android.AndroidInputLifecycle/extensions/android_frames_extension.ts
index 1328cbb..58717cc 100644
--- a/ui/src/plugins/com.android.AndroidInputLifecycle/extensions/android_frames_extension.ts
+++ b/ui/src/plugins/com.android.AndroidInputLifecycle/extensions/android_frames_extension.ts
@@ -52,6 +52,24 @@
         tsField: 'ts_do_frame',
         durField: 'dur_do_frame',
       },
+      {
+        key: 'draw_frames',
+        headerName: 'DrawFrames',
+        sequenceNumber: 6000,
+        idField: 'id_draw_frames',
+        trackField: 'track_draw_frames',
+        tsField: 'ts_draw_frames',
+        durField: 'dur_draw_frames',
+      },
+      {
+        key: 'sf',
+        headerName: 'SurfaceFlinger',
+        sequenceNumber: 7000,
+        idField: 'id_sf',
+        trackField: 'track_sf',
+        tsField: 'ts_sf',
+        durField: 'dur_sf',
+      },
     ];
   }
 
@@ -60,12 +78,25 @@
     sliceId: number,
   ): Promise<string | undefined> {
     const query = `
+      WITH selected_vsync AS (
+        -- Case 1: Selected Choreographer slice
+        SELECT chor.frame_id AS vsync_id, chor.upid
+        FROM android_frames_choreographer_do_frame chor
+        WHERE chor.id = ${sliceId}
+        UNION ALL
+        -- Case 2: Selected DrawFrames slice
+        SELECT d.frame_id AS vsync_id, d.upid
+        FROM android_frames_draw_frame d
+        WHERE d.id = ${sliceId}
+        UNION ALL
+        -- Case 3: Selected SurfaceFlinger composite slice
+        SELECT sf.app_vsync AS vsync_id, sf.app_upid AS upid
+        FROM _input_sf_resolved sf
+        WHERE sf.id = ${sliceId}
+      )
       SELECT e.input_event_id AS input_id
       FROM android_input_events e
-      JOIN android_frames_choreographer_do_frame chor 
-        ON chor.frame_id = e.frame_id
-        AND chor.upid = e.upid
-      WHERE chor.id = ${sliceId}
+      JOIN selected_vsync cv ON e.frame_id = cv.vsync_id AND e.upid = cv.upid
       LIMIT 1
     `;
     const result = await trace.engine.query(query);