ui/src/plugins/dev.perfetto.RestorePinnedTracks/index.ts - third_party/perfetto - Git at Google

 // Copyright (C) 2024 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 {TrackNode} from '../../public/workspace';
 import {Trace} from '../../public/trace';
 import {PerfettoPlugin} from '../../public/plugin';
 import {TrackDescriptor} from '../../public/track';
 import {z} from 'zod';
 import {
   base64Decode,
   base64Encode,
   utf8Decode,
   utf8Encode,
 } from '../../base/string_utils';
 import {assertExists} from '../../base/logging';

 const PLUGIN_ID = 'dev.perfetto.RestorePinnedTrack';
 const SAVED_TRACKS_KEY = `${PLUGIN_ID}#savedPerfettoTracks`;

 const RESTORE_COMMAND_ID = `${PLUGIN_ID}#restore`;

 /**
  * Fuzzy save and restore of pinned tracks.
  *
  * Tries to persist pinned tracks. Uses full string matching between track name
  * and group name. When no match is found for a saved track, it tries again
  * without numbers.
  */
 export default class implements PerfettoPlugin {
   static readonly id = PLUGIN_ID;
   private ctx!: Trace;

   async onTraceLoad(ctx: Trace): Promise<void> {
     this.ctx = ctx;

     ctx.commands.registerCommand({
       id: `${PLUGIN_ID}#save`,
       name: 'Save: Pinned tracks',
       callback: () => {
         this.savedState = {
           ...this.savedState,
           tracks: this.getCurrentPinnedTracks(),
         };
       },
     });
     ctx.commands.registerCommand({
       id: RESTORE_COMMAND_ID,
       name: 'Restore: Pinned tracks',
       callback: () => {
         const tracks = this.savedState?.tracks;
         if (!tracks) {
           alert('No saved tracks. Use the Save command first');
           return;
         }
         this.restoreTracks(tracks);
       },
     });

     ctx.commands.registerCommand({
       id: `${PLUGIN_ID}#saveByName`,
       name: 'Save by name: Pinned tracks',
       callback: async () => {
         const res = await this.ctx.omnibox.prompt(
           'Give a name to the pinned set of tracks',
         );
         if (res) {
           const rawTracksByName = this.savedState?.tracksByName ?? [];
           const tracksByNameMap = new Map(
             rawTracksByName.map((x) => [x.name, x.tracks]),
           );
           tracksByNameMap.set(
             base64Encode(utf8Encode(res)),
             this.getCurrentPinnedTracks(),
           );
           this.savedState = {
             ...this.savedState,
             tracksByName: Array.from(tracksByNameMap.entries()).map(
               ([k, v]) => ({
                 name: k,
                 tracks: v,
               }),
             ),
           };
         }
       },
     });
     ctx.commands.registerCommand({
       id: `${PLUGIN_ID}#restoreByName`,
       name: 'Restore by name: Pinned tracks',
       callback: async () => {
         const tracksByName = this.savedState?.tracksByName ?? [];
         if (tracksByName.length === 0) {
           alert('No saved tracks. Use the Save by name command first');
           return;
         }
         const res = await this.ctx.omnibox.prompt(
           'Select name of set of pinned tracks to restore',
           tracksByName.map((x) => ({
             key: x.name,
             displayName: utf8Decode(base64Decode(x.name)),
           })),
         );
         if (res) {
           const tracks = assertExists(
             tracksByName.find((x) => x.name === res)?.tracks,
           );
           this.restoreTracks(tracks);
         }
       },
     });
   }

   private restoreTracks(tracks: ReadonlyArray<SavedPinnedTrack>) {
     const localTracks = this.ctx.workspace.flatTracks.map((track) => ({
       savedTrack: this.toSavedTrack(track),
       track: track,
     }));
     tracks.forEach((trackToRestore) => {
       const foundTrack = this.findMatchingTrack(localTracks, trackToRestore);
       if (foundTrack) {
         foundTrack.pin();
       } else {
         console.warn(
           '[RestorePinnedTracks] No track found that matches',
           trackToRestore,
         );
       }
     });
   }

   private getCurrentPinnedTracks() {
     return this.ctx.workspace.pinnedTracks.map((track) =>
       this.toSavedTrack(track),
     );
   }

   private findMatchingTrack(
     localTracks: Array<LocalTrack>,
     savedTrack: SavedPinnedTrack,
   ): TrackNode | undefined {
     let mostSimilarTrack: LocalTrack | undefined = undefined;
     let mostSimilarTrackDifferenceScore: number = 0;

     for (let i = 0; i < localTracks.length; i++) {
       const localTrack = localTracks[i];
       const differenceScore = this.calculateSimilarityScore(
         localTrack.savedTrack,
         savedTrack,
       );

       // Return immediately if we found the exact match
       if (differenceScore === Number.MAX_SAFE_INTEGER) {
         return localTrack.track;
       }

       // Ignore too different objects
       if (differenceScore === 0) {
         continue;
       }

       if (differenceScore > mostSimilarTrackDifferenceScore) {
         mostSimilarTrackDifferenceScore = differenceScore;
         mostSimilarTrack = localTrack;
       }
     }

     return mostSimilarTrack?.track || undefined;
   }

   /**
    * Returns the similarity score where 0 means the objects are completely
    * different, and the higher the number, the smaller the difference is.
    * Returns Number.MAX_SAFE_INTEGER if the objects are completely equal.
    * We attempt a fuzzy match based on the similarity score.
    * For example, one of the ways we do this is we remove the numbers
    * from the title to potentially pin a "similar" track from a different trace.
    * Removing numbers allows flexibility; for instance, with multiple 'sysui'
    * processes (e.g. track group name: "com.android.systemui 123") without
    * this approach, any could be mistakenly pinned. The goal is to restore
    * specific tracks within the same trace, ensuring that a previously pinned
    * track is pinned again.
    * If the specific process with that PID is unavailable, pinning any
    * other process matching the package name is attempted.
    * @param track1 first saved track to compare
    * @param track2 second saved track to compare
    * @private
    */
   private calculateSimilarityScore(
     track1: SavedPinnedTrack,
     track2: SavedPinnedTrack,
   ): number {
     // Return immediately when objects are equal
     if (
       track1.trackName === track2.trackName &&
       track1.groupName === track2.groupName &&
       track1.pluginId === track2.pluginId &&
       track1.kind === track2.kind &&
       track1.isMainThread === track2.isMainThread
     ) {
       return Number.MAX_SAFE_INTEGER;
     }

     let similarityScore = 0;
     if (track1.trackName === track2.trackName) {
       similarityScore += 100;
     } else if (
       this.removeNumbers(track1.trackName) ===
       this.removeNumbers(track2.trackName)
     ) {
       similarityScore += 50;
     }

     if (track1.groupName === track2.groupName) {
       similarityScore += 90;
     } else if (
       this.removeNumbers(track1.groupName) ===
       this.removeNumbers(track2.groupName)
     ) {
       similarityScore += 45;
     }

     // Do not consider other parameters if there is no match in name/group
     if (similarityScore === 0) return similarityScore;

     if (track1.pluginId === track2.pluginId) {
       similarityScore += 30;
     }

     if (track1.kind === track2.kind) {
       similarityScore += 20;
     }

     if (track1.isMainThread === track2.isMainThread) {
       similarityScore += 10;
     }

     return similarityScore;
   }

   private removeNumbers(inputString?: string): string | undefined {
     return inputString?.replace(/\d+/g, '');
   }

   private toSavedTrack(track: TrackNode): SavedPinnedTrack {
     let trackDescriptor: TrackDescriptor | undefined = undefined;
     if (track.uri != undefined) {
       trackDescriptor = this.ctx.tracks.getTrack(track.uri);
     }

     return {
       groupName: groupName(track),
       trackName: track.title,
       pluginId: trackDescriptor?.pluginId,
       kind: trackDescriptor?.tags?.kind,
       isMainThread: trackDescriptor?.chips?.includes('main thread') || false,
     };
   }

   private get savedState(): SavedState | undefined {
     const savedStateString = window.localStorage.getItem(SAVED_TRACKS_KEY);
     if (!savedStateString) {
       return undefined;
     }
     const savedState = SAVED_STATE_SCHEMA.safeParse(
       JSON.parse(savedStateString),
     );
     if (!savedState.success) {
       return undefined;
     }
     return savedState.data;
   }

   private set savedState(state: SavedState) {
     window.localStorage.setItem(SAVED_TRACKS_KEY, JSON.stringify(state));
   }
 }

 // Return the displayname of the containing group
 // If the track is a child of a workspace, return undefined...
 function groupName(track: TrackNode): string | undefined {
   const parent = track.parent;
   if (parent instanceof TrackNode) {
     return parent.title;
   }
   return undefined;
 }

 const SAVED_PINNED_TRACK_SCHEMA = z
   .object({
     // Optional: group name for the track. Usually matches with process name.
     groupName: z.string().optional(),
     // Track name to restore.
     trackName: z.string(),
     // Plugin used to create this track
     pluginId: z.string().optional(),
     // Kind of the track
     kind: z.string().optional(),
     // If it's a thread track, it should be true in case it's a main thread track
     isMainThread: z.boolean(),
   })
   .readonly();

 type SavedPinnedTrack = z.infer<typeof SAVED_PINNED_TRACK_SCHEMA>;

 const SAVED_STATE_SCHEMA = z
   .object({
     tracks: z.array(SAVED_PINNED_TRACK_SCHEMA).optional().readonly(),
     tracksByName: z
       .array(
         z
           .object({
             name: z.string(),
             tracks: z.array(SAVED_PINNED_TRACK_SCHEMA).readonly(),
           })
           .readonly(),
       )
       .optional()
       .readonly(),
   })
   .readonly();

 type SavedState = z.infer<typeof SAVED_STATE_SCHEMA>;

 interface LocalTrack {
   readonly savedTrack: SavedPinnedTrack;
   readonly track: TrackNode;
 }
	// Copyright (C) 2024 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 {TrackNode} from '../../public/workspace';
	import {Trace} from '../../public/trace';
	import {PerfettoPlugin} from '../../public/plugin';
	import {TrackDescriptor} from '../../public/track';
	import {z} from 'zod';
	import {
	base64Decode,
	base64Encode,
	utf8Decode,
	utf8Encode,
	} from '../../base/string_utils';
	import {assertExists} from '../../base/logging';

	const PLUGIN_ID = 'dev.perfetto.RestorePinnedTrack';
	const SAVED_TRACKS_KEY = `${PLUGIN_ID}#savedPerfettoTracks`;

	const RESTORE_COMMAND_ID = `${PLUGIN_ID}#restore`;

	/**
	* Fuzzy save and restore of pinned tracks.
	*
	* Tries to persist pinned tracks. Uses full string matching between track name
	* and group name. When no match is found for a saved track, it tries again
	* without numbers.
	*/
	export default class implements PerfettoPlugin {
	static readonly id = PLUGIN_ID;
	private ctx!: Trace;

	async onTraceLoad(ctx: Trace): Promise<void> {
	this.ctx = ctx;

	ctx.commands.registerCommand({
	id: `${PLUGIN_ID}#save`,
	name: 'Save: Pinned tracks',
	callback: () => {
	this.savedState = {
	...this.savedState,
	tracks: this.getCurrentPinnedTracks(),
	};
	},
	});
	ctx.commands.registerCommand({
	id: RESTORE_COMMAND_ID,
	name: 'Restore: Pinned tracks',
	callback: () => {
	const tracks = this.savedState?.tracks;
	if (!tracks) {
	alert('No saved tracks. Use the Save command first');
	return;
	}
	this.restoreTracks(tracks);
	},
	});

	ctx.commands.registerCommand({
	id: `${PLUGIN_ID}#saveByName`,
	name: 'Save by name: Pinned tracks',
	callback: async () => {
	const res = await this.ctx.omnibox.prompt(
	'Give a name to the pinned set of tracks',
	);
	if (res) {
	const rawTracksByName = this.savedState?.tracksByName ?? [];
	const tracksByNameMap = new Map(
	rawTracksByName.map((x) => [x.name, x.tracks]),
	);
	tracksByNameMap.set(
	base64Encode(utf8Encode(res)),
	this.getCurrentPinnedTracks(),
	);
	this.savedState = {
	...this.savedState,
	tracksByName: Array.from(tracksByNameMap.entries()).map(
	([k, v]) => ({
	name: k,
	tracks: v,
	}),
	),
	};
	}
	},
	});
	ctx.commands.registerCommand({
	id: `${PLUGIN_ID}#restoreByName`,
	name: 'Restore by name: Pinned tracks',
	callback: async () => {
	const tracksByName = this.savedState?.tracksByName ?? [];
	if (tracksByName.length === 0) {
	alert('No saved tracks. Use the Save by name command first');
	return;
	}
	const res = await this.ctx.omnibox.prompt(
	'Select name of set of pinned tracks to restore',
	tracksByName.map((x) => ({
	key: x.name,
	displayName: utf8Decode(base64Decode(x.name)),
	})),
	);
	if (res) {
	const tracks = assertExists(
	tracksByName.find((x) => x.name === res)?.tracks,
	);
	this.restoreTracks(tracks);
	}
	},
	});
	}

	private restoreTracks(tracks: ReadonlyArray<SavedPinnedTrack>) {
	const localTracks = this.ctx.workspace.flatTracks.map((track) => ({
	savedTrack: this.toSavedTrack(track),
	track: track,
	}));
	tracks.forEach((trackToRestore) => {
	const foundTrack = this.findMatchingTrack(localTracks, trackToRestore);
	if (foundTrack) {
	foundTrack.pin();
	} else {
	console.warn(
	'[RestorePinnedTracks] No track found that matches',
	trackToRestore,
	);
	}
	});
	}

	private getCurrentPinnedTracks() {
	return this.ctx.workspace.pinnedTracks.map((track) =>
	this.toSavedTrack(track),
	);
	}

	private findMatchingTrack(
	localTracks: Array<LocalTrack>,
	savedTrack: SavedPinnedTrack,
	): TrackNode \| undefined {
	let mostSimilarTrack: LocalTrack \| undefined = undefined;
	let mostSimilarTrackDifferenceScore: number = 0;

	for (let i = 0; i < localTracks.length; i++) {
	const localTrack = localTracks[i];
	const differenceScore = this.calculateSimilarityScore(
	localTrack.savedTrack,
	savedTrack,
	);

	// Return immediately if we found the exact match
	if (differenceScore === Number.MAX_SAFE_INTEGER) {
	return localTrack.track;
	}

	// Ignore too different objects
	if (differenceScore === 0) {
	continue;
	}

	if (differenceScore > mostSimilarTrackDifferenceScore) {
	mostSimilarTrackDifferenceScore = differenceScore;
	mostSimilarTrack = localTrack;
	}
	}

	return mostSimilarTrack?.track \|\| undefined;
	}

	/**
	* Returns the similarity score where 0 means the objects are completely
	* different, and the higher the number, the smaller the difference is.
	* Returns Number.MAX_SAFE_INTEGER if the objects are completely equal.
	* We attempt a fuzzy match based on the similarity score.
	* For example, one of the ways we do this is we remove the numbers
	* from the title to potentially pin a "similar" track from a different trace.
	* Removing numbers allows flexibility; for instance, with multiple 'sysui'
	* processes (e.g. track group name: "com.android.systemui 123") without
	* this approach, any could be mistakenly pinned. The goal is to restore
	* specific tracks within the same trace, ensuring that a previously pinned
	* track is pinned again.
	* If the specific process with that PID is unavailable, pinning any
	* other process matching the package name is attempted.
	* @param track1 first saved track to compare
	* @param track2 second saved track to compare
	* @private
	*/
	private calculateSimilarityScore(
	track1: SavedPinnedTrack,
	track2: SavedPinnedTrack,
	): number {
	// Return immediately when objects are equal
	if (
	track1.trackName === track2.trackName &&
	track1.groupName === track2.groupName &&
	track1.pluginId === track2.pluginId &&
	track1.kind === track2.kind &&
	track1.isMainThread === track2.isMainThread
	) {
	return Number.MAX_SAFE_INTEGER;
	}

	let similarityScore = 0;
	if (track1.trackName === track2.trackName) {
	similarityScore += 100;
	} else if (
	this.removeNumbers(track1.trackName) ===
	this.removeNumbers(track2.trackName)
	) {
	similarityScore += 50;
	}

	if (track1.groupName === track2.groupName) {
	similarityScore += 90;
	} else if (
	this.removeNumbers(track1.groupName) ===
	this.removeNumbers(track2.groupName)
	) {
	similarityScore += 45;
	}

	// Do not consider other parameters if there is no match in name/group
	if (similarityScore === 0) return similarityScore;

	if (track1.pluginId === track2.pluginId) {
	similarityScore += 30;
	}

	if (track1.kind === track2.kind) {
	similarityScore += 20;
	}

	if (track1.isMainThread === track2.isMainThread) {
	similarityScore += 10;
	}

	return similarityScore;
	}

	private removeNumbers(inputString?: string): string \| undefined {
	return inputString?.replace(/\d+/g, '');
	}

	private toSavedTrack(track: TrackNode): SavedPinnedTrack {
	let trackDescriptor: TrackDescriptor \| undefined = undefined;
	if (track.uri != undefined) {
	trackDescriptor = this.ctx.tracks.getTrack(track.uri);
	}

	return {
	groupName: groupName(track),
	trackName: track.title,
	pluginId: trackDescriptor?.pluginId,
	kind: trackDescriptor?.tags?.kind,
	isMainThread: trackDescriptor?.chips?.includes('main thread') \|\| false,
	};
	}

	private get savedState(): SavedState \| undefined {
	const savedStateString = window.localStorage.getItem(SAVED_TRACKS_KEY);
	if (!savedStateString) {
	return undefined;
	}
	const savedState = SAVED_STATE_SCHEMA.safeParse(
	JSON.parse(savedStateString),
	);
	if (!savedState.success) {
	return undefined;
	}
	return savedState.data;
	}

	private set savedState(state: SavedState) {
	window.localStorage.setItem(SAVED_TRACKS_KEY, JSON.stringify(state));
	}
	}

	// Return the displayname of the containing group
	// If the track is a child of a workspace, return undefined...
	function groupName(track: TrackNode): string \| undefined {
	const parent = track.parent;
	if (parent instanceof TrackNode) {
	return parent.title;
	}
	return undefined;
	}

	const SAVED_PINNED_TRACK_SCHEMA = z
	.object({
	// Optional: group name for the track. Usually matches with process name.
	groupName: z.string().optional(),
	// Track name to restore.
	trackName: z.string(),
	// Plugin used to create this track
	pluginId: z.string().optional(),
	// Kind of the track
	kind: z.string().optional(),
	// If it's a thread track, it should be true in case it's a main thread track
	isMainThread: z.boolean(),
	})
	.readonly();

	type SavedPinnedTrack = z.infer<typeof SAVED_PINNED_TRACK_SCHEMA>;

	const SAVED_STATE_SCHEMA = z
	.object({
	tracks: z.array(SAVED_PINNED_TRACK_SCHEMA).optional().readonly(),
	tracksByName: z
	.array(
	z
	.object({
	name: z.string(),
	tracks: z.array(SAVED_PINNED_TRACK_SCHEMA).readonly(),
	})
	.readonly(),
	)
	.optional()
	.readonly(),
	})
	.readonly();

	type SavedState = z.infer<typeof SAVED_STATE_SCHEMA>;

	interface LocalTrack {
	readonly savedTrack: SavedPinnedTrack;
	readonly track: TrackNode;
	}