ui/src/plugins/com.android.DayExplorer/index.ts

// Copyright (C) 2023 The Android Open Source Project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

import {uuidv4} from '../../base/uuid';
import {Trace} from '../../public/trace';
import StandardGroupsPlugin from '../dev.perfetto.StandardGroups';
import {PerfettoPlugin} from '../../public/plugin';
import {TrackNode} from '../../public/workspace';
import {STR, LONG, NUM, LONG_NULL} from '../../trace_processor/query_result';
import {SourceDataset} from '../../trace_processor/dataset';
import {AreaSelection, areaSelectionsEqual} from '../../public/selection';
import {Flamegraph} from '../../widgets/flamegraph';
import {
  metricsFromTableOrSubquery,
  QueryFlamegraph,
} from '../../components/query_flamegraph';
import SupportPlugin from '../com.android.AndroidLongBatterySupport';
import {
  CounterRowSchema,
  CounterTrack,
} from '../../components/tracks/counter_track';

const DAY_EXPLORER_TRACK_KIND = 'day_explorer_counter_track';

export default class implements PerfettoPlugin {
  static readonly id = 'com.android.DayExplorer';
  static readonly dependencies = [StandardGroupsPlugin, SupportPlugin];

  private support(ctx: Trace) {
    return ctx.plugins.getPlugin(SupportPlugin);
  }

  async addDayExplorerCounters(
    ctx: Trace,
    support: SupportPlugin,
    groupName: string,
    limit: number,
  ): Promise<void> {
    await ctx.engine.query(
      `INCLUDE PERFETTO MODULE
          google3.wireless.android.telemetry.trace_extractor.modules.day_explorer.perfetto_ui_blames`,
    );

    const group = support.getOrCreateGroup(ctx, groupName);
    await this.addDayExplorerRecursive(ctx, group, limit, -1n);
  }

  private async addDayExplorerRecursive(
    ctx: Trace,
    parent: TrackNode,
    limit: number,
    parentId: bigint,
  ): Promise<void> {
    const children = await ctx.engine.query(`
      SELECT track_id, display_name, cast(round(total_energy_uws / 3600000) as int) as energy_mwh
      FROM day_explorer_ui_hierarchy
      WHERE (${parentId} >= 0 AND parent_id = ${parentId})
         OR (${parentId} < 0 AND parent_id IS NULL)
      ORDER BY energy_mwh DESC
      LIMIT ${limit}
    `);

    const childIter = children.iter({
      track_id: LONG,
      display_name: STR,
      energy_mwh: LONG,
    });

    for (; childIter.valid(); childIter.next()) {
      const dataset = new SourceDataset({
        src: `
          SELECT
            ts,
            power_mw AS value,
            track_id
          FROM day_explorer_ui_hierarchy_per_ts
        `,
        schema: {
          ts: LONG,
          value: NUM,
        },
        filter: {
          col: 'track_id',
          eq: childIter.track_id,
        },
      });
      const groupKey = `_day_explorer_ui_hierarchy_under_${parentId}`;
      const trackName = `${childIter.display_name} - ${childIter.energy_mwh}mWh`;
      const node = await this.createDayExplorerTrack(
        ctx,
        trackName,
        groupKey,
        dataset,
      );
      parent.addChildInOrder(node);
      await this.addDayExplorerRecursive(ctx, node, limit, childIter.track_id);
    }
  }

  private async createDayExplorerTrack(
    ctx: Trace,
    name: string,
    groupKey: string,
    dataset: SourceDataset<CounterRowSchema>,
  ): Promise<TrackNode> {
    const uri = `/day_explorer_${uuidv4()}`;
    const renderer = await CounterTrack.createMaterialized({
      trace: ctx,
      uri,
      dataset,
      defaultOptions: {
        yRangeSharingKey: groupKey,
      },
    });

    ctx.tracks.registerTrack({
      uri,
      renderer,
      tags: {
        kinds: [DAY_EXPLORER_TRACK_KIND],
      },
    });

    return new TrackNode({
      name,
      uri,
    });
  }

  private createDayExplorerFlameGraphPanel(trace: Trace) {
    let previousSelection: AreaSelection | undefined;
    let flamegraph: QueryFlamegraph | undefined;
    return {
      id: 'day_explorer_flamegraph_selection',
      name: 'Day Explorer Flamegraph',
      render: (selection: AreaSelection) => {
        const selectionChanged =
          previousSelection === undefined ||
          !areaSelectionsEqual(previousSelection, selection);
        previousSelection = selection;
        if (selectionChanged) {
          flamegraph = this.computeDayExplorerFlameGraph(trace, selection);
        }

        if (flamegraph === undefined) {
          return undefined;
        }

        return {isLoading: false, content: flamegraph.render()};
      },
    };
  }

  computeDayExplorerFlameGraph(trace: Trace, currentSelection: AreaSelection) {
    // The flame graph will be shown when any day explorer track is in the area
    // selection. The selection is used to filter by time, but not by track. All
    // day explorer tracks are considered for the graph.
    let hasDayExplorer = false;
    for (const trackInfo of currentSelection.tracks) {
      if (trackInfo?.tags?.kinds?.includes(DAY_EXPLORER_TRACK_KIND)) {
        hasDayExplorer = true;
        break;
      }
    }
    if (!hasDayExplorer) {
      return undefined;
    }
    const metrics = metricsFromTableOrSubquery(
      `
        (
          WITH
            total_energy AS (
              SELECT track_id, parent_id, display_name, SUM(energy_uws) AS energy_uws
              FROM day_explorer_ui_hierarchy_per_ts
              WHERE ts >= ${currentSelection.start}
                AND ts <= ${currentSelection.end}
              GROUP BY 1, 2, 3
            ),
            with_child AS (
              SELECT
                *,
                (
                  SELECT IFNULL(SUM(energy_uws), 0)
                  FROM total_energy
                  WHERE parent_id = P.track_id
                ) AS child_energy
              FROM total_energy AS P
            )
          SELECT
            track_id AS id,
            parent_id AS parentId,
            display_name AS name,
            cast(round((energy_uws - child_energy) / 1000) as int) AS self_count
          FROM with_child
        )
      `,
      [
        {
          name: 'Energy mWs',
          unit: '',
          columnName: 'self_count',
        },
      ],
    );
    return new QueryFlamegraph(trace, metrics, {
      state: Flamegraph.createDefaultState(metrics),
    });
  }

  async addDayExplorerBehaviors(
    ctx: Trace,
    support: SupportPlugin,
    groupName: string,
  ): Promise<void> {
    const e = ctx.engine;

    await e.query(
      `INCLUDE PERFETTO MODULE
          google3.wireless.android.telemetry.trace_extractor.modules.day_explorer.perfetto_ui_blames`,
    );

    await support.addSliceTrack(
      ctx,
      'Day Explorer Behaviors',
      new SourceDataset({
        src: `
          SELECT
            ts,
            dur,
            behavior as name
          FROM day_explorer_behaviors
        `,
        schema: {
          ts: LONG,
          dur: LONG_NULL,
          name: STR,
        },
      }),
      groupName,
      false,
    );
  }

  async addDayExplorerCommand(
    ctx: Trace,
    support: SupportPlugin,
    features: Set<string>,
  ): Promise<void> {
    if (features.has('google3')) {
      ctx.commands.registerCommand({
        id: 'com.android.DayExplorerBlamesByCategory',
        name: 'Add tracks: Day Explorer',
        callback: async () => {
          const limitStr = await ctx.omnibox.prompt(
            'Maximum results per group',
          );
          const limit = Number(limitStr);
          if (!isFinite(limit) || limit <= 0) {
            alert('Positive number required');
            return;
          }
          await this.addDayExplorerBehaviors(ctx, support, 'Day explorer');
          await this.addDayExplorerCounters(
            ctx,
            support,
            'Day explorer',
            limit,
          );
        },
      });
    }
  }

  async onTraceLoad(ctx: Trace): Promise<void> {
    const support = this.support(ctx);
    const features = await support.features(ctx.engine);

    ctx.selection.registerAreaSelectionTab(
      this.createDayExplorerFlameGraphPanel(ctx),
    );

    if (features.has('google3')) {
      await this.addDayExplorerCommand(ctx, support, features);
    }
  }
}