Update cardToPinnedCopyCache when resolving unresolved pins. (#5837)

Googlers, see b/240338417 "Clicking the refresh UI button unpins pinned graphs".

When a user loads a TensorBoard with pinned cards in the URL then clicks on the TensorBoard reload button, the  pinned cards disappear!

Explanation:
* Soon after page load, when handing tag metadata, unresolved pins are resolved into actual pinned cards. Unfortunately, the cardToPinnedCopyCache property is not provided with the information:
  * https://cs.opensource.google/tensorflow/tensorboard/+/master:tensorboard/webapp/metrics/store/metrics_reducers.ts;l=565;drc=2f5661647c264a3cd5758a04a7de0beb0c604747
* When user then presses the TensorBoard reload button and then tag metadata is reloaded, the cache is considered empty and the actual set of pinned cards is also cleared:
  * https://cs.opensource.google/tensorflow/tensorboard/+/master:tensorboard/webapp/metrics/store/metrics_reducers.ts;l=550;drc=2f5661647c264a3cd5758a04a7de0beb0c604747

The solution is to ensure that cardToPinnedCopyCache is also populated when resolving unresolved pins.

tensorboard/webapp/metrics/store/metrics_reducers.ts

@@ -372,6 +372,7 @@ const reducer = createReducer(
       state.cardList,
       state.cardMetadataMap,
       state.cardToPinnedCopy,
+      state.cardToPinnedCopyCache,
       state.pinnedCardToOriginal,
       state.cardStepIndex
     );
@@ -396,7 +397,6 @@ const reducer = createReducer(
     const newState = {
       ...state,
       ...resolvedResult,
-      cardToPinnedCopyCache: resolvedResult.cardToPinnedCopy,
       settingOverrides: newSettings,
     };
 
@@ -567,6 +567,7 @@ const reducer = createReducer(
         nextCardList,
         nextCardMetadataMap,
         nextCardToPinnedCopy,
+        state.cardToPinnedCopyCache,
         nextPinnedCardToOriginal,
         nextCardStepIndex
       );
@@ -922,12 +923,13 @@ const reducer = createReducer(
         const resolvedResult = buildOrReturnStateWithPinnedCopy(
           cardId,
           nextCardToPinnedCopy,
+          nextCardToPinnedCopyCache,
           nextPinnedCardToOriginal,
           nextCardStepIndexMap,
           nextCardMetadataMap
         );
         nextCardToPinnedCopy = resolvedResult.cardToPinnedCopy;
-        nextCardToPinnedCopyCache = resolvedResult.cardToPinnedCopy;
+        nextCardToPinnedCopyCache = resolvedResult.cardToPinnedCopyCache;
         nextPinnedCardToOriginal = resolvedResult.pinnedCardToOriginal;
         nextCardMetadataMap = resolvedResult.cardMetadataMap;
         nextCardStepIndexMap = resolvedResult.cardStepIndex;

tensorboard/webapp/metrics/store/metrics_reducers_test.ts

@@ -715,6 +715,7 @@ describe('metrics reducers', () => {
         cardList,
         cardStepIndex,
         cardToPinnedCopy,
+        cardToPinnedCopyCache,
         pinnedCardToOriginal,
         unresolvedImportedPinnedCards,
       } = nextState;
@@ -723,6 +724,7 @@ describe('metrics reducers', () => {
         cardList,
         cardStepIndex,
         cardToPinnedCopy,
+        cardToPinnedCopyCache,
         pinnedCardToOriginal,
         unresolvedImportedPinnedCards,
       }).toEqual({
@@ -738,6 +740,9 @@ describe('metrics reducers', () => {
           }),
         },
         cardToPinnedCopy: new Map([[expectedCardId, expectedPinnedCopyId]]),
+        cardToPinnedCopyCache: new Map([
+          [expectedCardId, expectedPinnedCopyId],
+        ]),
         pinnedCardToOriginal: new Map([[expectedPinnedCopyId, expectedCardId]]),
         unresolvedImportedPinnedCards: [
           {plugin: PluginType.SCALARS, tag: 'tagB'},

tensorboard/webapp/metrics/store/metrics_store_internal_utils.ts

@@ -42,6 +42,7 @@ type ResolvedPinPartialState = Pick<
   MetricsState,
   | 'cardMetadataMap'
   | 'cardToPinnedCopy'
+  | 'cardToPinnedCopyCache'
   | 'pinnedCardToOriginal'
   | 'cardStepIndex'
 >;
@@ -203,6 +204,7 @@ export function buildOrReturnStateWithUnresolvedImportedPins(
   nonPinnedCards: NonPinnedCardId[],
   cardMetadataMap: CardMetadataMap,
   cardToPinnedCopy: CardToPinnedCard,
+  cardToPinnedCopyCache: CardToPinnedCard,
   pinnedCardToOriginal: PinnedCardToCard,
   cardStepIndexMap: CardStepIndexMap
 ): ResolvedPinPartialState & {unresolvedImportedPinnedCards: CardUniqueInfo[]} {
@@ -224,13 +226,15 @@ export function buildOrReturnStateWithUnresolvedImportedPins(
       unresolvedImportedPinnedCards,
       cardMetadataMap,
       cardToPinnedCopy,
+      cardToPinnedCopyCache,
       pinnedCardToOriginal,
       cardStepIndex: cardStepIndexMap,
     };
   }
 
   let stateWithResolvedPins = {
     cardToPinnedCopy,
+    cardToPinnedCopyCache,
     pinnedCardToOriginal,
     cardStepIndex: cardStepIndexMap,
     cardMetadataMap,
@@ -239,6 +243,7 @@ export function buildOrReturnStateWithUnresolvedImportedPins(
     stateWithResolvedPins = buildOrReturnStateWithPinnedCopy(
       cardToPin,
       stateWithResolvedPins.cardToPinnedCopy,
+      stateWithResolvedPins.cardToPinnedCopyCache,
       stateWithResolvedPins.pinnedCardToOriginal,
       stateWithResolvedPins.cardStepIndex,
       stateWithResolvedPins.cardMetadataMap
@@ -258,6 +263,7 @@ export function buildOrReturnStateWithUnresolvedImportedPins(
 export function buildOrReturnStateWithPinnedCopy(
   cardId: NonPinnedCardId,
   cardToPinnedCopy: CardToPinnedCard,
+  cardToPinnedCopyCache: CardToPinnedCard,
   pinnedCardToOriginal: PinnedCardToCard,
   cardStepIndexMap: CardStepIndexMap,
   cardMetadataMap: CardMetadataMap
@@ -266,20 +272,23 @@ export function buildOrReturnStateWithPinnedCopy(
   if (cardToPinnedCopy.has(cardId)) {
     return {
       cardToPinnedCopy,
+      cardToPinnedCopyCache,
       pinnedCardToOriginal,
       cardStepIndex: cardStepIndexMap,
       cardMetadataMap,
     };
   }
 
   const nextCardToPinnedCopy = new Map(cardToPinnedCopy);
+  const nextCardToPinnedCopyCache = new Map(cardToPinnedCopyCache);
   const nextPinnedCardToOriginal = new Map(pinnedCardToOriginal);
   const nextCardStepIndexMap = {...cardStepIndexMap};
   const nextCardMetadataMap = {...cardMetadataMap};
 
   // Create a pinned copy. Copies step index from the original card.
   const pinnedCardId = getPinnedCardId(cardId);
   nextCardToPinnedCopy.set(cardId, pinnedCardId);
+  nextCardToPinnedCopyCache.set(cardId, pinnedCardId);
   nextPinnedCardToOriginal.set(pinnedCardId, cardId);
   if (cardStepIndexMap.hasOwnProperty(cardId)) {
     nextCardStepIndexMap[pinnedCardId] = cardStepIndexMap[cardId];
@@ -293,6 +302,7 @@ export function buildOrReturnStateWithPinnedCopy(
 
   return {
     cardToPinnedCopy: nextCardToPinnedCopy,
+    cardToPinnedCopyCache: nextCardToPinnedCopyCache,
     pinnedCardToOriginal: nextPinnedCardToOriginal,
     cardStepIndex: nextCardStepIndexMap,
     cardMetadataMap: nextCardMetadataMap,

tensorboard/webapp/metrics/store/metrics_store_internal_utils_test.ts

@@ -373,6 +373,7 @@ describe('metrics store utils', () => {
         {card1: {plugin: PluginType.SCALARS, tag: 'accuracy', runId: null}},
         new Map(),
         new Map(),
+        new Map(),
         {card1: buildStepIndexMetadata({index: 2})}
       );
 
@@ -381,6 +382,9 @@ describe('metrics store utils', () => {
       expect(result.cardToPinnedCopy).toEqual(
         new Map([['card1', pinnedCardId]])
       );
+      expect(result.cardToPinnedCopyCache).toEqual(
+        new Map([['card1', pinnedCardId]])
+      );
       expect(result.pinnedCardToOriginal).toEqual(
         new Map([[pinnedCardId, 'card1']])
       );
@@ -391,19 +395,22 @@ describe('metrics store utils', () => {
     it('adds a pinned copy properly', () => {
       const {
         cardToPinnedCopy,
+        cardToPinnedCopyCache,
         pinnedCardToOriginal,
         cardStepIndex,
         cardMetadataMap,
       } = buildOrReturnStateWithPinnedCopy(
         'card1',
         new Map(),
         new Map(),
+        new Map(),
         {card1: buildStepIndexMetadata({index: 2})},
         {card1: createCardMetadata()}
       );
       const pinnedCardId = getPinnedCardId('card1');
 
       expect(cardToPinnedCopy).toEqual(new Map([['card1', pinnedCardId]]));
+      expect(cardToPinnedCopyCache).toEqual(new Map([['card1', pinnedCardId]]));
       expect(pinnedCardToOriginal).toEqual(new Map([[pinnedCardId, 'card1']]));
       expect(cardStepIndex).toEqual({
         card1: buildStepIndexMetadata({index: 2}),
@@ -417,17 +424,26 @@ describe('metrics store utils', () => {
 
     it('throws if the original card does not have metadata', () => {
       expect(() => {
-        buildOrReturnStateWithPinnedCopy('card1', new Map(), new Map(), {}, {});
+        buildOrReturnStateWithPinnedCopy(
+          'card1',
+          new Map(),
+          new Map(),
+          new Map(),
+          {},
+          {}
+        );
       }).toThrow();
     });
 
     it('no-ops if the card already has a pinned copy', () => {
       const cardToPinnedCopy = new Map([['card1', 'card-pin1']]);
+      const cardToPinnedCopyCache = new Map([['card1', 'card-pin1']]);
       const pinnedCardToOriginal = new Map([['card-pin1', 'card1']]);
       const cardStepIndexMap = {};
       const cardMetadataMap = {card1: createCardMetadata()};
       const originals = {
         cardToPinnedCopy: new Map(cardToPinnedCopy),
+        cardToPinnedCopyCache: new Map(cardToPinnedCopyCache),
         pinnedCardToOriginal: new Map(pinnedCardToOriginal),
         cardStepIndexMap: {...cardStepIndexMap},
         cardMetadataMap: {...cardMetadataMap},
@@ -436,12 +452,16 @@ describe('metrics store utils', () => {
       const result = buildOrReturnStateWithPinnedCopy(
         'card1',
         cardToPinnedCopy,
+        cardToPinnedCopyCache,
         pinnedCardToOriginal,
         cardStepIndexMap,
         cardMetadataMap
       );
 
       expect(result.cardToPinnedCopy).toEqual(originals.cardToPinnedCopy);
+      expect(result.cardToPinnedCopyCache).toEqual(
+        originals.cardToPinnedCopyCache
+      );
       expect(result.pinnedCardToOriginal).toEqual(
         originals.pinnedCardToOriginal
       );