ui: Use dataset for the thread state aggregation tab

Bug: 383066495
Change-Id: Ia523ef982c02d47944ffbb80965cc00c76fd2c97

ui/src/plugins/dev.perfetto.ThreadState/thread_state_selection_aggregator.ts

@@ -12,32 +12,27 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import {exists} from '../../base/utils';
 import {ColumnDef, Sorting, ThreadStateExtra} from '../../public/aggregation';
 import {AreaSelection} from '../../public/selection';
-import {THREAD_STATE_TRACK_KIND} from '../../public/track_kinds';
 import {Engine} from '../../trace_processor/engine';
-import {NUM, NUM_NULL, STR_NULL} from '../../trace_processor/query_result';
+import {
+  LONG,
+  NUM,
+  NUM_NULL,
+  STR,
+  STR_NULL,
+} from '../../trace_processor/query_result';
 import {AreaSelectionAggregator} from '../../public/selection';
+import {UnionDataset} from '../../trace_processor/dataset';
 import {translateState} from '../../components/sql_utils/thread_state';
 import {TrackDescriptor} from '../../public/track';
 
 export class ThreadStateSelectionAggregator implements AreaSelectionAggregator {
   readonly id = 'thread_state_aggregation';
-  private utids?: number[];
-
-  setThreadStateUtids(tracks: ReadonlyArray<TrackDescriptor>) {
-    this.utids = [];
-    for (const trackInfo of tracks) {
-      if (trackInfo?.tags?.kind === THREAD_STATE_TRACK_KIND) {
-        exists(trackInfo.tags.utid) && this.utids.push(trackInfo.tags.utid);
-      }
-    }
-  }
 
   async createAggregateView(engine: Engine, area: AreaSelection) {
-    this.setThreadStateUtids(area.tracks);
-    if (this.utids === undefined || this.utids.length === 0) return false;
+    const dataset = this.getDatasetFromTracks(area.tracks);
+    if (dataset === undefined) return false;
 
     await engine.query(`
       create or replace perfetto table ${this.id} as
@@ -50,35 +45,34 @@ export class ThreadStateSelectionAggregator implements AreaSelectionAggregator {
         sum(tstate.dur) AS total_dur,
         sum(tstate.dur) / count() as avg_dur,
         count() as occurrences
-      from thread_state tstate
+      from (${dataset.query()}) tstate
       join thread using (utid)
       left join process using (upid)
       where
-        utid in (${this.utids})
-        and ts + dur > ${area.start}
+        ts + dur > ${area.start}
         and ts < ${area.end}
       group by utid, concat_state
     `);
+
     return true;
   }
 
   async getExtra(
     engine: Engine,
     area: AreaSelection,
   ): Promise<ThreadStateExtra | void> {
-    this.setThreadStateUtids(area.tracks);
-    if (this.utids === undefined || this.utids.length === 0) return;
+    const dataset = this.getDatasetFromTracks(area.tracks);
+    if (dataset === undefined) return;
 
     const query = `
       select
         state,
         io_wait as ioWait,
         sum(dur) as totalDur
-      from thread
-      join thread_state using (utid)
-      where utid in (${this.utids})
-        and thread_state.ts + thread_state.dur > ${area.start}
-        and thread_state.ts < ${area.end}
+      from (${dataset.query()}) tstate
+      join thread using (utid)
+      where tstate.ts + tstate.dur > ${area.start}
+        and tstate.ts < ${area.end}
       group by state, io_wait
     `;
     const result = await engine.query(query);
@@ -170,4 +164,24 @@ export class ThreadStateSelectionAggregator implements AreaSelectionAggregator {
   getDefaultSorting(): Sorting {
     return {column: 'total_dur', direction: 'DESC'};
   }
+
+  // Creates an optimized dataset containing the thread state events within a
+  // given list of tracks, or returns undefined if no compatible tracks are
+  // present in the list.
+  private getDatasetFromTracks(tracks: ReadonlyArray<TrackDescriptor>) {
+    const desiredSchema = {
+      dur: LONG,
+      io_wait: NUM_NULL,
+      state: STR,
+      utid: NUM,
+    };
+    const validDatasets = tracks
+      .map((track) => track.track.getDataset?.())
+      .filter((ds) => ds !== undefined)
+      .filter((ds) => ds.implements(desiredSchema));
+    if (validDatasets.length === 0) {
+      return undefined;
+    }
+    return new UnionDataset(validDatasets).optimize();
+  }
 }

ui/src/plugins/dev.perfetto.ThreadState/thread_state_track.ts

@@ -22,12 +22,13 @@ import {
   SLICE_LAYOUT_FLAT_DEFAULTS,
   SliceLayout,
 } from '../../components/tracks/slice_layout';
-import {NUM_NULL, STR} from '../../trace_processor/query_result';
+import {LONG, NUM, NUM_NULL, STR} from '../../trace_processor/query_result';
 import {Slice} from '../../public/track';
 import {translateState} from '../../components/sql_utils/thread_state';
 import {TrackEventDetails, TrackEventSelection} from '../../public/selection';
 import {ThreadStateDetailsPanel} from './thread_state_details_panel';
 import {Trace} from '../../public/trace';
+import {Dataset, SourceDataset} from '../../trace_processor/dataset';
 
 export const THREAD_STATE_ROW = {
   ...BASE_ROW,
@@ -71,6 +72,25 @@ export class ThreadStateTrack extends BaseSliceTrack<Slice, ThreadStateRow> {
     `;
   }
 
+  getDataset(): Dataset | undefined {
+    return new SourceDataset({
+      src: 'thread_state',
+      schema: {
+        id: NUM,
+        ts: LONG,
+        dur: LONG,
+        cpu: NUM,
+        state: STR,
+        io_wait: NUM_NULL,
+        utid: NUM,
+      },
+      filter: {
+        col: 'utid',
+        eq: this.utid,
+      },
+    });
+  }
+
   rowToSlice(row: ThreadStateRow): Slice {
     const baseSlice = this.rowToSliceBase(row);
     const ioWait = row.ioWait === null ? undefined : !!row.ioWait;