time-namespaced state: Improve how we decide when to rehydrate state …

…from URL. (tensorflow#5549)

There are some browser history navigations where we were not appropriately rehydrating URL state. The root cause was that we were only rehydrating URL state at most one time for each namespace. But, in reality, we might need to do it multiple times: one time for the experiment list view and another time for the dashboard view.

(Googlers, see internal bug b/216829997).

This change fixes the issue by adding another dimension when determining whether we can rehydrate URL state. The dimension is based on `RouteKind` but actually represented by a new `DeepLinkGroup` enum.

The mapping of `RouteKind` to `DeepLinkGroup` is explicitly defined in the app routing code:
* `DeepLinkGroup.EXPERIMENTS` = `[RouteKind.EXPERIMENTS]`
* `DeepLinkGroup.DASHBOARD` = `[RouteKind.EXPERIMENT, RouteKind.COMPARE_EXPERIMENT]`

I considered making `DeepLinkGroup` (or some equivalent) a per-app configuration, either in the route configurations or in the `DeepLinkProvider` interface. But after playing around with these ideas, I realized the API and code would be quite messy for little gain. Besides, we don't currently see a world where the `DeepLinkGroup` definition differs between different versions of Tensorboard.

I also renamed `knownNamespaceIds` to `rehydatedDeepLinks` to make the purpose clearer and narrower.

tensorboard/webapp/app_routing/effects/app_routing_effects.ts

@@ -42,6 +42,7 @@ import {DirtyUpdatesRegistryModule} from '../dirty_updates_registry_module';
 import {
   arePathsAndQueryParamsEqual,
   areSameRouteKindAndExperiments,
+  canRehydrateDeepLink,
   getRouteNamespaceId,
   serializeCompareExperimentParams,
 } from '../internal_utils';
@@ -52,7 +53,7 @@ import {RouteRegistryModule} from '../route_registry_module';
 import {
   getActiveNamespaceId,
   getActiveRoute,
-  getKnownNamespaceIds,
+  getRehydratedDeepLinks,
 } from '../store/app_routing_selectors';
 import {Route, RouteKind, RouteParams, SerializableQueryParams} from '../types';
 
@@ -412,24 +413,25 @@ export class AppRoutingEffects {
           })
         );
       }),
-      withLatestFrom(this.store.select(getKnownNamespaceIds)),
-      tap(([{routeMatch, options}, knownNamespaceIds]) => {
+      withLatestFrom(this.store.select(getRehydratedDeepLinks)),
+      tap(([{routeMatch, options}, rehydratedDeepLinks]) => {
         // Possibly rehydrate state from the URL.
-        //
-        // We do this on "browser initiated" events (like a page load or
-        // navigation in browser history) but we only do this when we don't yet
-        // have in-memory state for the namespace being navigated to.
 
-        const isKnownNamespace =
-          options.namespaceUpdate.option ===
-            NamespaceUpdateOption.FROM_HISTORY &&
-          knownNamespaceIds.has(options.namespaceUpdate.namespaceId);
+        if (!options.browserInitiated || !routeMatch.deepLinkProvider) {
+          return;
+        }
 
         if (
-          !options.browserInitiated ||
-          isKnownNamespace ||
-          !routeMatch.deepLinkProvider
+          options.namespaceUpdate.option ===
+            NamespaceUpdateOption.FROM_HISTORY &&
+          !canRehydrateDeepLink(
+            routeMatch.routeKind,
+            options.namespaceUpdate.namespaceId,
+            rehydratedDeepLinks
+          )
         ) {
+          // A deeplink has already been rehydrated for this RouteKind/Namespace
+          // combination so don't do it again.
           return;
         }
 

tensorboard/webapp/app_routing/effects/app_routing_effects_test.ts

@@ -35,10 +35,11 @@ import {RouteRegistryModule} from '../route_registry_module';
 import {
   getActiveNamespaceId,
   getActiveRoute,
-  getKnownNamespaceIds,
+  getRehydratedDeepLinks,
 } from '../store/app_routing_selectors';
 import {buildRoute, provideLocationTesting, TestableLocation} from '../testing';
 import {
+  DeepLinkGroup,
   DirtyUpdates,
   Navigation,
   NavigationFromHistory,
@@ -227,7 +228,7 @@ describe('app_routing_effects', () => {
 
     store.overrideSelector(getActiveRoute, null);
     store.overrideSelector(getActiveNamespaceId, null);
-    store.overrideSelector(getKnownNamespaceIds, new Set<string>());
+    store.overrideSelector(getRehydratedDeepLinks, []);
     store.overrideSelector(getEnabledTimeNamespacedState, false);
     actualActions = [];
 
@@ -858,7 +859,7 @@ describe('app_routing_effects', () => {
 
     describe('deeplink reads', () => {
       beforeEach(() => {
-        store.overrideSelector(getKnownNamespaceIds, new Set<string>());
+        store.overrideSelector(getRehydratedDeepLinks, []);
         store.refreshState();
       });
 
@@ -915,14 +916,20 @@ describe('app_routing_effects', () => {
         }));
       });
 
-      it('dispatches stateRehydratedFromUrl when unknown namespace id', fakeAsync(() => {
+      it('dispatches stateRehydratedFromUrl if canRehydrateDeepLink()', fakeAsync(() => {
         deserializeQueryParamsSpy.and.returnValue({a: 'A', b: 'B'});
         getPathSpy.and.returnValue('/compare/a:b');
 
-        store.overrideSelector(
-          getKnownNamespaceIds,
-          new Set(['namespaceA', 'namespaceC'])
-        );
+        store.overrideSelector(getRehydratedDeepLinks, [
+          {
+            namespaceId: 'namespaceA',
+            deepLinkGroup: DeepLinkGroup.DASHBOARD,
+          },
+          {
+            namespaceId: 'namespaceB',
+            deepLinkGroup: DeepLinkGroup.EXPERIMENTS,
+          },
+        ]);
         store.refreshState();
 
         onPopStateSubject.next({
@@ -946,14 +953,20 @@ describe('app_routing_effects', () => {
         ]);
       }));
 
-      it('does not dispatch stateRehydratedFromUrl when known namespace id', fakeAsync(() => {
+      it('does not dispatch stateRehydratedFromUrl if not canRehydrateDeepLink()', fakeAsync(() => {
         deserializeQueryParamsSpy.and.returnValue({a: 'A', b: 'B'});
         getPathSpy.and.returnValue('/compare/a:b');
 
-        store.overrideSelector(
-          getKnownNamespaceIds,
-          new Set(['namespaceA', 'namespaceB', 'namespaceC'])
-        );
+        store.overrideSelector(getRehydratedDeepLinks, [
+          {
+            namespaceId: 'namespaceA',
+            deepLinkGroup: DeepLinkGroup.DASHBOARD,
+          },
+          {
+            namespaceId: 'namespaceB',
+            deepLinkGroup: DeepLinkGroup.DASHBOARD,
+          },
+        ]);
         store.refreshState();
 
         onPopStateSubject.next({

tensorboard/webapp/app_routing/internal_utils.ts

@@ -14,8 +14,10 @@ limitations under the License.
 ==============================================================================*/
 import {
   CompareRouteParams,
+  DeepLinkGroup,
   DEFAULT_EXPERIMENT_ID,
   ExperimentRouteParams,
+  RehydratedDeepLink,
   Route,
   RouteKind,
   RouteParams,
@@ -190,3 +192,55 @@ export function arePathsAndQueryParamsEqual(
     return paramA.key === paramB.key && paramA.value === paramB.value;
   });
 }
+
+/**
+ * Maps route kinds to their deep link groups.
+ */
+export function getDeepLinkGroup(routeKind: RouteKind): DeepLinkGroup | null {
+  switch (routeKind) {
+    case RouteKind.EXPERIMENTS:
+      return DeepLinkGroup.EXPERIMENTS;
+    case RouteKind.EXPERIMENT:
+    case RouteKind.COMPARE_EXPERIMENT:
+      return DeepLinkGroup.DASHBOARD;
+    case RouteKind.UNKNOWN:
+    case RouteKind.NOT_SET:
+      return null;
+  }
+}
+
+/**
+ * Determines whether a combination of route kind and namespace id should have
+ * a URL deep link rehydrated into state.
+ *
+ * We limit the number of times we rehydrate deep links as a user navigates
+ * through browser history using back and forward buttons -- rehydrating only
+ * one time for each combination of namespaceId and deepLinkGroup.
+ *
+ * So, for example, when:
+ *
+ * 1. User reloads page into experiment list, rehydrate state for the
+ *    EXPERIMENTS deep link group.
+ * 2. User then navigates back to a compare view, rehydrate state for the
+ *    DASHBOARD deep link group.
+ * 3. User then navigates back to an experiment view, DO NOT rehydrate state
+ *    again for the DASHBOARD deep link group as it was rehydrated in step (2).
+ * 4. If user then navigates back to an experiment view but in a different
+ *    namespace, then once again rehydrate state for the DASHBOARD deep link
+ *    group.
+ */
+export function canRehydrateDeepLink(
+  routeKind: RouteKind,
+  namespaceId: string,
+  rehydratedDeepLinks: RehydratedDeepLink[]
+) {
+  const deepLinkGroup = getDeepLinkGroup(routeKind);
+  return (
+    deepLinkGroup !== null &&
+    !rehydratedDeepLinks.some(
+      (rehydratedDeepLink) =>
+        rehydratedDeepLink.deepLinkGroup === deepLinkGroup &&
+        rehydratedDeepLink.namespaceId === namespaceId
+    )
+  );
+}

tensorboard/webapp/app_routing/internal_utils_test.ts

@@ -19,7 +19,7 @@ import {
   buildExperimentRouteFromId,
   buildRoute,
 } from './testing';
-import {RouteKind} from './types';
+import {DeepLinkGroup, RouteKind} from './types';
 
 describe('app_routing/utils', () => {
   describe('#parseCompareExperimentStr', () => {
@@ -123,7 +123,7 @@ describe('app_routing/utils', () => {
     });
   });
 
-  describe('getRouteNamespaceId', () => {
+  describe('#getRouteNamespaceId', () => {
     [
       {
         kind: RouteKind.COMPARE_EXPERIMENT,
@@ -175,7 +175,7 @@ describe('app_routing/utils', () => {
     });
   });
 
-  describe('areSameRouteKindAndExperiments', () => {
+  describe('#areSameRouteKindAndExperiments', () => {
     it('returns true when both routes are null', () => {
       expect(utils.areSameRouteKindAndExperiments(null, null)).toBeTrue();
     });
@@ -406,4 +406,82 @@ describe('app_routing/utils', () => {
       ).toBe(false);
     });
   });
+
+  describe('#getDeepLinkGroup', () => {
+    it('maps RouteKind to DeepLinkGroup', () => {
+      expect(utils.getDeepLinkGroup(RouteKind.EXPERIMENTS)).toEqual(
+        DeepLinkGroup.EXPERIMENTS
+      );
+      expect(utils.getDeepLinkGroup(RouteKind.EXPERIMENT)).toEqual(
+        DeepLinkGroup.DASHBOARD
+      );
+      expect(utils.getDeepLinkGroup(RouteKind.COMPARE_EXPERIMENT)).toEqual(
+        DeepLinkGroup.DASHBOARD
+      );
+      expect(utils.getDeepLinkGroup(RouteKind.UNKNOWN)).toBeNull();
+      expect(utils.getDeepLinkGroup(RouteKind.NOT_SET)).toBeNull();
+    });
+  });
+
+  describe('#canRehydrateDeepLink', () => {
+    it('allows rehydration if namespaceId does not match', () => {
+      expect(
+        utils.canRehydrateDeepLink(RouteKind.EXPERIMENTS, 'namespaceC', [
+          {
+            deepLinkGroup: DeepLinkGroup.EXPERIMENTS,
+            namespaceId: 'namespaceA',
+          },
+          {
+            deepLinkGroup: DeepLinkGroup.EXPERIMENTS,
+            namespaceId: 'namespaceB',
+          },
+        ])
+      ).toBeTrue();
+    });
+
+    it('allows rehydration if deepLinkGroup does not match', () => {
+      expect(
+        utils.canRehydrateDeepLink(RouteKind.COMPARE_EXPERIMENT, 'namespaceA', [
+          {
+            deepLinkGroup: DeepLinkGroup.EXPERIMENTS,
+            namespaceId: 'namespaceA',
+          },
+          {
+            deepLinkGroup: DeepLinkGroup.EXPERIMENTS,
+            namespaceId: 'namespaceB',
+          },
+        ])
+      ).toBeTrue();
+    });
+
+    it('does not allow rehydration if match is found', () => {
+      expect(
+        utils.canRehydrateDeepLink(RouteKind.EXPERIMENTS, 'namespaceA', [
+          {
+            deepLinkGroup: DeepLinkGroup.EXPERIMENTS,
+            namespaceId: 'namespaceA',
+          },
+          {
+            deepLinkGroup: DeepLinkGroup.EXPERIMENTS,
+            namespaceId: 'namespaceB',
+          },
+        ])
+      ).toBeFalse();
+    });
+
+    it('does not allow rehydration if route kind has null deep link group', () => {
+      expect(
+        utils.canRehydrateDeepLink(RouteKind.UNKNOWN, 'namespaceA', [
+          {
+            deepLinkGroup: DeepLinkGroup.EXPERIMENTS,
+            namespaceId: 'namespaceA',
+          },
+          {
+            deepLinkGroup: DeepLinkGroup.EXPERIMENTS,
+            namespaceId: 'namespaceB',
+          },
+        ])
+      ).toBeFalse();
+    });
+  });
 });

tensorboard/webapp/app_routing/store/app_routing_reducers.ts

@@ -14,13 +14,14 @@ limitations under the License.
 ==============================================================================*/
 import {Action, createReducer, on} from '@ngrx/store';
 import * as actions from '../actions';
+import {canRehydrateDeepLink, getDeepLinkGroup} from '../internal_utils';
 import {AppRoutingState} from './app_routing_types';
 
 const initialState: AppRoutingState = {
   activeRoute: null,
   nextRoute: null,
   activeNamespaceId: null,
-  knownNamespaceIds: new Set(),
+  rehydratedDeepLinks: [],
   registeredRouteKeys: new Set(),
 };
 
@@ -30,17 +31,28 @@ const reducer = createReducer(
     return {...state, nextRoute: after};
   }),
   on(actions.navigated, (state, {after, afterNamespaceId}) => {
-    let knownNamespaceIds = state.knownNamespaceIds;
-    if (!state.knownNamespaceIds.has(afterNamespaceId)) {
-      knownNamespaceIds = new Set(state.knownNamespaceIds);
-      knownNamespaceIds.add(afterNamespaceId);
+    let rehydratedDeepLinks = state.rehydratedDeepLinks;
+    if (
+      canRehydrateDeepLink(
+        after.routeKind,
+        afterNamespaceId,
+        rehydratedDeepLinks
+      )
+    ) {
+      rehydratedDeepLinks = [...rehydratedDeepLinks];
+      rehydratedDeepLinks.push({
+        // Note: getDeepLinkGroup() should return non-null given that
+        // canRehydrateDeepLink() returned true.
+        deepLinkGroup: getDeepLinkGroup(after.routeKind)!,
+        namespaceId: afterNamespaceId,
+      });
     }
     return {
       ...state,
       activeRoute: after,
       nextRoute: null,
       activeNamespaceId: afterNamespaceId,
-      knownNamespaceIds,
+      rehydratedDeepLinks,
     };
   }),
   on(actions.routeConfigLoaded, (state, {routeKinds}) => {