Translate URL track-index state through profile sanitization

src/actions/publish.ts

@@ -17,17 +17,32 @@ import {
   getDataSource,
   getProfileNameForStorage,
   getUrlPredictor,
+  getGlobalTrackOrder,
+  getHiddenGlobalTracks,
+  getHiddenLocalTracksByPid,
+  getLocalTrackOrderByPid,
+  getTabFilter,
 } from 'firefox-profiler/selectors/url-state';
 import {
   getProfile,
   getZeroAt,
   getCommittedRange,
+  getGlobalTracks,
+  getLocalTracksByPid,
+  getTabToThreadIndexesMap,
 } from 'firefox-profiler/selectors/profile';
 import { viewProfile } from './receive-profile';
 import { ensureExists } from 'firefox-profiler/utils/types';
 import { extractProfileTokenFromJwt } from 'firefox-profiler/utils/jwt';
 import { withHistoryReplaceStateSync } from 'firefox-profiler/app-logic/url-handling';
 import { persistUploadedProfileInformationToDb } from 'firefox-profiler/app-logic/uploaded-profiles-db';
+import {
+  computeGlobalTracks,
+  computeLocalTracksByPid,
+  computeOldTrackIndexToNewTrackIndexMap,
+  computeHiddenTracksAfterSanitization,
+  computeTrackOrderAfterSanitization,
+} from 'firefox-profiler/profile-logic/tracks';
 
 import type {
   Action,
@@ -36,6 +51,11 @@ import type {
   StartEndRange,
   State,
   Profile,
+  Pid,
+  RawCounter,
+  CounterIndex,
+  ThreadIndex,
+  TrackIndex,
   ProfileIndexTranslationMaps,
 } from 'firefox-profiler/types';
 import { compress } from 'firefox-profiler/utils/gz';
@@ -232,6 +252,47 @@ async function persistJustUploadedProfileInformationToDb(
   }
 }
 
+/**
+ * Map each old CounterIndex to its new CounterIndex across a sanitization
+ * step. Sanitization removes counters whose parent thread is removed; the
+ * survivors keep their relative order. Each counter is identified by
+ * (pid, category, name, mainThreadIndex), with the old-side mainThreadIndex
+ * normalized through `oldThreadIndexToNew`.
+ */
+function _computeOldCounterIndexToNew(
+  oldCounters: RawCounter[] | null | undefined,
+  newCounters: RawCounter[] | null | undefined,
+  oldThreadIndexToNew: Map<ThreadIndex, ThreadIndex>
+): Map<CounterIndex, CounterIndex> {
+  const result = new Map<CounterIndex, CounterIndex>();
+  if (!oldCounters || !newCounters) {
+    return result;
+  }
+  const newKeyToIndex = new Map<string, CounterIndex>();
+  for (let i = 0; i < newCounters.length; i++) {
+    const c = newCounters[i];
+    newKeyToIndex.set(
+      `${c.pid}|${c.category}|${c.name}|${c.mainThreadIndex}`,
+      i
+    );
+  }
+  for (let i = 0; i < oldCounters.length; i++) {
+    const c = oldCounters[i];
+    const newMainThreadIndex = oldThreadIndexToNew.get(c.mainThreadIndex);
+    if (newMainThreadIndex === undefined) {
+      // The counter's parent thread is gone in the new profile, so the
+      // counter is gone too.
+      continue;
+    }
+    const key = `${c.pid}|${c.category}|${c.name}|${newMainThreadIndex}`;
+    const newIndex = newKeyToIndex.get(key);
+    if (newIndex !== undefined) {
+      result.set(i, newIndex);
+    }
+  }
+  return result;
+}
+
 export type ProfileEncodingResult =
   | {
       type: 'SUCCESS';
@@ -413,6 +474,95 @@ export function attemptToPublish(
       if (removeProfileInformation) {
         const { committedRanges, translationMaps, profile } =
           sanitizedInformation;
+
+        // When sanitization re-indexed tracks, translate the URL state's
+        // track-index references through old → new track-index space so the
+        // user's choices for tracks that survived sanitization stay attached
+        // to the right tracks.
+        let remappedHiddenGlobalTracks: Set<TrackIndex> | null = null;
+        let remappedGlobalTrackOrder: TrackIndex[] | null = null;
+        let remappedHiddenLocalTracksByPid: Map<Pid, Set<TrackIndex>> | null =
+          null;
+        let remappedLocalTrackOrderByPid: Map<Pid, TrackIndex[]> | null = null;
+
+        const oldThreadIndexToNew = translationMaps?.oldThreadIndexToNew;
+        if (oldThreadIndexToNew) {
+          const oldProfile = getProfile(prePublishedState);
+          const oldCounterIndexToNew = _computeOldCounterIndexToNew(
+            oldProfile.counters,
+            profile.counters,
+            oldThreadIndexToNew
+          );
+
+          const tabToThreadIndexesMap =
+            getTabToThreadIndexesMap(prePublishedState);
+          const tabFilter = getTabFilter(prePublishedState);
+          const newGlobalTracks = computeGlobalTracks(
+            profile,
+            tabFilter,
+            tabToThreadIndexesMap
+          );
+          const newLocalTracksByPid = computeLocalTracksByPid(
+            profile,
+            newGlobalTracks
+          );
+
+          const oldGlobalTracks = getGlobalTracks(prePublishedState);
+          const oldLocalTracksByPid = getLocalTracksByPid(prePublishedState);
+
+          const globalOldToNew = computeOldTrackIndexToNewTrackIndexMap({
+            oldTracks: oldGlobalTracks,
+            newTracks: newGlobalTracks,
+            oldThreadIndexToNew,
+            oldCounterIndexToNew,
+          });
+
+          remappedHiddenGlobalTracks = computeHiddenTracksAfterSanitization({
+            oldHiddenTracks: getHiddenGlobalTracks(prePublishedState),
+            oldTrackIndexToNewTrackIndex: globalOldToNew,
+          });
+          remappedGlobalTrackOrder = computeTrackOrderAfterSanitization({
+            oldTrackOrder: getGlobalTrackOrder(prePublishedState),
+            oldTrackIndexToNewTrackIndex: globalOldToNew,
+          });
+
+          remappedHiddenLocalTracksByPid = new Map();
+          remappedLocalTrackOrderByPid = new Map();
+          const oldHiddenLocalByPid =
+            getHiddenLocalTracksByPid(prePublishedState);
+          const oldLocalOrderByPid = getLocalTrackOrderByPid(prePublishedState);
+
+          for (const [pid, newLocalTracks] of newLocalTracksByPid) {
+            const oldLocalTracks = oldLocalTracksByPid.get(pid) ?? [];
+            const localOldToNew = computeOldTrackIndexToNewTrackIndexMap({
+              oldTracks: oldLocalTracks,
+              newTracks: newLocalTracks,
+              oldThreadIndexToNew,
+              oldCounterIndexToNew,
+            });
+            const oldHiddenForPid = oldHiddenLocalByPid.get(pid);
+            if (oldHiddenForPid !== undefined) {
+              remappedHiddenLocalTracksByPid.set(
+                pid,
+                computeHiddenTracksAfterSanitization({
+                  oldHiddenTracks: oldHiddenForPid,
+                  oldTrackIndexToNewTrackIndex: localOldToNew,
+                })
+              );
+            }
+            const oldOrderForPid = oldLocalOrderByPid.get(pid);
+            if (oldOrderForPid !== undefined) {
+              remappedLocalTrackOrderByPid.set(
+                pid,
+                computeTrackOrderAfterSanitization({
+                  oldTrackOrder: oldOrderForPid,
+                  oldTrackIndexToNewTrackIndex: localOldToNew,
+                })
+              );
+            }
+          }
+        }
+
         // Hide the old UI gracefully.
         await dispatch(hideStaleProfile());
 
@@ -423,7 +573,11 @@ export function attemptToPublish(
             committedRanges,
             translationMaps,
             profileName,
-            prePublishedState
+            prePublishedState,
+            remappedHiddenGlobalTracks,
+            remappedGlobalTrackOrder,
+            remappedHiddenLocalTracksByPid,
+            remappedLocalTrackOrderByPid
           )
         );
 
@@ -504,13 +658,22 @@ export function resetUploadState(): Action {
 /**
  * Report to the UrlState that the profile was sanitized. This will re-map any stored
  * indexes or information that has been sanitized away.
+ *
+ * The four track-index payload fields carry URL state translated into the
+ * post-sanitization track-index space when sanitization re-indexed tracks
+ * (translationMaps.oldThreadIndexToNew is non-null). They are null when no
+ * remap is needed; in that case the existing reducer state is left untouched.
  */
 export function profileSanitized(
   hash: string,
   committedRanges: StartEndRange[] | null,
   translationMaps: ProfileIndexTranslationMaps | null,
   profileName: string,
-  prePublishedState: State | null
+  prePublishedState: State | null,
+  hiddenGlobalTracks: Set<TrackIndex> | null = null,
+  globalTrackOrder: TrackIndex[] | null = null,
+  hiddenLocalTracksByPid: Map<Pid, Set<TrackIndex>> | null = null,
+  localTrackOrderByPid: Map<Pid, TrackIndex[]> | null = null
 ): Action {
   return {
     type: 'SANITIZED_PROFILE_PUBLISHED',
@@ -519,6 +682,10 @@ export function profileSanitized(
     translationMaps,
     profileName,
     prePublishedState,
+    hiddenGlobalTracks,
+    globalTrackOrder,
+    hiddenLocalTracksByPid,
+    localTrackOrderByPid,
   };
 }
 

src/profile-logic/tracks.ts

@@ -10,8 +10,10 @@ import type {
   Pid,
   GlobalTrack,
   LocalTrack,
+  Track,
   TrackIndex,
   RawCounter,
+  CounterIndex,
   Tid,
   TrackReference,
   TabID,
@@ -496,6 +498,148 @@ export function addProcessCPUTracksForProcess(
   return newLocalTracksByPid;
 }
 
+/**
+ * Build an identity key for a track based on properties that survive
+ * sanitization. When a translation map is provided for the dimension a track
+ * keys on, the source-side index is normalized through it; a missing entry
+ * means the track was removed and the key is null. Marker tracks key on a
+ * string-table index that sanitization reshuffles, so they're not matched
+ * here.
+ */
+function _trackIdentityKey(
+  track: Track,
+  threadIndexTranslation: Map<ThreadIndex, ThreadIndex> | null,
+  counterIndexTranslation: Map<CounterIndex, CounterIndex> | null
+): string | null {
+  switch (track.type) {
+    case 'process':
+      return `process:${track.pid}`;
+    case 'screenshots':
+      return `screenshots:${track.id}`;
+    case 'visual-progress':
+    case 'perceptual-visual-progress':
+    case 'contentful-visual-progress':
+      return track.type;
+    case 'thread':
+    case 'network':
+    case 'ipc':
+    case 'event-delay': {
+      let ti: ThreadIndex | undefined = track.threadIndex;
+      if (threadIndexTranslation !== null) {
+        ti = threadIndexTranslation.get(track.threadIndex);
+        if (ti === undefined) {
+          return null;
+        }
+      }
+      return `${track.type}:${ti}`;
+    }
+    case 'counter': {
+      let ci: CounterIndex | undefined = track.counterIndex;
+      if (counterIndexTranslation !== null) {
+        ci = counterIndexTranslation.get(track.counterIndex);
+        if (ci === undefined) {
+          return null;
+        }
+      }
+      return `counter:${ci}`;
+    }
+    case 'marker':
+      return null;
+    default:
+      throw assertExhaustiveCheck(track, 'Unhandled Track type.');
+  }
+}
+
+/**
+ * Map each old TrackIndex to its new TrackIndex when both old and new track
+ * lists describe the same profile across a sanitization step. Tracks are
+ * matched by stable identity (pid, screenshot id, threadIndex, counterIndex,
+ * or visual-progress singleton); old-side thread- and counter-index values
+ * are normalized through the supplied translation maps before comparison.
+ * Tracks with no match in `newTracks` are absent from the result.
+ */
+export function computeOldTrackIndexToNewTrackIndexMap({
+  oldTracks,
+  newTracks,
+  oldThreadIndexToNew,
+  oldCounterIndexToNew,
+}: {
+  readonly oldTracks: ReadonlyArray<Track>;
+  readonly newTracks: ReadonlyArray<Track>;
+  readonly oldThreadIndexToNew: Map<ThreadIndex, ThreadIndex> | null;
+  readonly oldCounterIndexToNew: Map<CounterIndex, CounterIndex> | null;
+}): Map<TrackIndex, TrackIndex> {
+  const newKeyToTrackIndex = new Map<string, TrackIndex>();
+  for (let i = 0; i < newTracks.length; i++) {
+    const key = _trackIdentityKey(newTracks[i], null, null);
+    if (key !== null) {
+      newKeyToTrackIndex.set(key, i);
+    }
+  }
+
+  const oldTrackIndexToNewTrackIndex = new Map<TrackIndex, TrackIndex>();
+  for (let i = 0; i < oldTracks.length; i++) {
+    const key = _trackIdentityKey(
+      oldTracks[i],
+      oldThreadIndexToNew,
+      oldCounterIndexToNew
+    );
+    if (key === null) {
+      continue;
+    }
+    const newIndex = newKeyToTrackIndex.get(key);
+    if (newIndex !== undefined) {
+      oldTrackIndexToNewTrackIndex.set(i, newIndex);
+    }
+  }
+
+  return oldTrackIndexToNewTrackIndex;
+}
+
+/**
+ * Translate a Set of hidden track indexes from old to new track-index space.
+ * Hidden tracks whose old index has no match in the new track list are dropped.
+ */
+export function computeHiddenTracksAfterSanitization({
+  oldHiddenTracks,
+  oldTrackIndexToNewTrackIndex,
+}: {
+  readonly oldHiddenTracks: ReadonlySet<TrackIndex>;
+  readonly oldTrackIndexToNewTrackIndex: Map<TrackIndex, TrackIndex>;
+}): Set<TrackIndex> {
+  const newHiddenTracks = new Set<TrackIndex>();
+  for (const oldIndex of oldHiddenTracks) {
+    const newIndex = oldTrackIndexToNewTrackIndex.get(oldIndex);
+    if (newIndex !== undefined) {
+      newHiddenTracks.add(newIndex);
+    }
+  }
+  return newHiddenTracks;
+}
+
+/**
+ * Translate a track-order array from old to new track-index space, preserving
+ * relative order. Indexes with no match in the new track list are dropped;
+ * `initializeGlobalTrackOrder` / `initializeLocalTrackOrderByPid` append any
+ * new tracks not represented in the input order.
+ */
+export function computeTrackOrderAfterSanitization({
+  oldTrackOrder,
+  oldTrackIndexToNewTrackIndex,
+}: {
+  readonly oldTrackOrder: ReadonlyArray<TrackIndex>;
+  readonly oldTrackIndexToNewTrackIndex: Map<TrackIndex, TrackIndex>;
+}): TrackIndex[] {
+  const newTrackOrder: TrackIndex[] = [];
+  for (const oldIndex of oldTrackOrder) {
+    const newIndex = oldTrackIndexToNewTrackIndex.get(oldIndex);
+    if (newIndex !== undefined) {
+      newTrackOrder.push(newIndex);
+    }
+  }
+  return newTrackOrder;
+}
+
 /**
  * Take a profile and figure out what GlobalTracks it contains.
  */