Add mechanism to include feature flags in Polymer HTTP requests (#5727)

This adds a mechanism that dumps the client-side FeatureFlag object into each HTTP request that originates from the TB Polymer code base.

(Googlers, see: http://go/tb-client-feature-flags-in-data-provider for the motivation and high level design.)

See #5718 for the corresponding change for Angular and #5721 for server-side support.

The data is dumped into a custom header. The format is effectively:
X-TensorBoard-Feature-Flags: JSON.stringify(currentFeatureFlags)

We implement this directly in the RequestManager object, through which all our Polymer HTTP requests are routed.

Additional things worth noting:

* I added a const for the header name to be shared between the HttpInterceptor and the RequestManager.

* I added a new angular<->polymer interop object, tf-feature-flags, modeled after some of the other interop objects we've added in the past. This interop object should now be the source of truth for accessing feature flags from the polymer code base. It is populated by feature_flags_effects soon after we fetch feature flag overrides from the data source.

* This will send ALL feature flags in the request header rather than a curated subset of them, just like in #5718. In a subsequent PR, we'll add support to mark certain feature flags as "sendToServer" and only send that subset.

tensorboard/components/BUILD

@@ -30,6 +30,7 @@ tf_ts_library(
     deps = [
         "//tensorboard/components/tf_backend",
         "//tensorboard/components/tf_color_scale",
+        "//tensorboard/components/tf_feature_flags",
         "//tensorboard/components/tf_globals",
         "//tensorboard/components/tf_markdown_view",
         "//tensorboard/components/tf_paginated_view",

tensorboard/components/polymer3_interop_helper.ts

@@ -18,6 +18,7 @@ limitations under the License.
  */
 import './tf_backend/tf-backend-polymer';
 import './tf_color_scale/tf-color-scale-polymer';
+import './tf_feature_flags/tf-feature-flags-polymer';
 import './tf_globals/globals-polymer';
 import './tf_paginated_view/tf-paginated-view-store';
 import './tf_storage/tf-storage-polymer';

tensorboard/components/tf_backend/BUILD

@@ -22,7 +22,9 @@ tf_ts_library(
     strict_checks = False,
     deps = [
         ":type",
+        "//tensorboard/components/tf_feature_flags",
         "//tensorboard/components/vz_sorting",
+        "//tensorboard/webapp/feature_flag/http:const",
         "@npm//@polymer/decorators",
         "@npm//@polymer/polymer",
         "@npm//@types/lodash",

tensorboard/components/tf_backend/requestManager.ts

@@ -12,6 +12,9 @@ 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 {FEATURE_FLAGS_HEADER_NAME} from '../../webapp/feature_flag/http/const';
+import {getFeatureFlags} from '../tf_feature_flags/feature-flags';
+
 interface ResolveReject {
   resolve: (x: unknown) => void;
   reject: (x: unknown) => void;
@@ -240,6 +243,10 @@ export class RequestManager {
         requestOptions.withCredentials,
         requestOptions.contentType
       );
+      req.setRequestHeader(
+        FEATURE_FLAGS_HEADER_NAME,
+        JSON.stringify(getFeatureFlags())
+      );
       req.onload = function () {
         if (req.status === 200) {
           resolve(JSON.parse(req.responseText) as any);

tensorboard/components/tf_feature_flags/BUILD

@@ -0,0 +1,38 @@
+load("//tensorboard/defs:defs.bzl", "tf_ng_web_test_suite", "tf_ts_library")
+
+package(default_visibility = ["//tensorboard:internal"])
+
+licenses(["notice"])
+
+tf_ts_library(
+    name = "tf_feature_flags",
+    srcs = [
+        "feature-flags.ts",
+        "tf-feature-flags-polymer.ts",
+    ],
+    deps = [
+        "//tensorboard/webapp/feature_flag:types",
+        "@npm//@polymer/decorators",
+        "@npm//@polymer/polymer",
+    ],
+)
+
+tf_ts_library(
+    name = "test_lib",
+    testonly = True,
+    srcs = [
+        "feature-flags-test.ts",
+    ],
+    deps = [
+        ":tf_feature_flags",
+        "//tensorboard/webapp/feature_flag:testing",
+        "@npm//@types/jasmine",
+    ],
+)
+
+tf_ng_web_test_suite(
+    name = "karma_test",
+    deps = [
+        ":test_lib",
+    ],
+)

tensorboard/components/tf_feature_flags/feature-flags-test.ts

@@ -0,0 +1,36 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+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 {buildFeatureFlag} from '../../webapp/feature_flag/testing';
+import * as feature_flags from './feature-flags';
+
+describe('feature-flags', () => {
+  beforeEach(() => {
+    feature_flags.initializeFeatureFlags();
+  });
+
+  it('sets and gets FeatureFlags', () => {
+    feature_flags.setFeatureFlags(buildFeatureFlag({inColab: true}));
+    expect(feature_flags.getFeatureFlags()).toEqual(
+      buildFeatureFlag({inColab: true})
+    );
+  });
+
+  it('throws Error if getFeatureFlags is called before setFeatureFlags', () => {
+    expect(() => feature_flags.getFeatureFlags()).toThrow(
+      new Error('FeatureFlags have not yet been determined by TensorBoard.')
+    );
+  });
+});

tensorboard/components/tf_feature_flags/feature-flags.ts

@@ -0,0 +1,50 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+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 {FeatureFlags} from '../../webapp/feature_flag/types';
+
+// FeatureFlags are managed and injected by the Angular portion of the TB code
+// base. In practice they are set one time soon after application load and never
+// again for the lifetime of the application.
+let _featureFlags: FeatureFlags | null;
+initializeFeatureFlags();
+
+export function initializeFeatureFlags(): void {
+  _featureFlags = null;
+}
+
+/**
+ * Sets the FeatureFlags for use in the Polymer portion of the TB code base.
+ *
+ * In practice this should only be called by the Angular portion of the TB code
+ * base immediately after it determines the final set of FeatureFlags.
+ */
+export function setFeatureFlags(featureFlags: FeatureFlags): void {
+  _featureFlags = featureFlags;
+}
+
+/**
+ * Retrieves FeatureFlags.
+ *
+ * @throws Error if FeatureFlags have not yet been set. In practice they should
+ *     be set soon after application load before any Polymer code is invoked.
+ *     This runtime check acts as a sanity check to enforce that assumption.
+ */
+export function getFeatureFlags(): FeatureFlags {
+  if (_featureFlags === null) {
+    throw Error('FeatureFlags have not yet been determined by TensorBoard.');
+  }
+  return _featureFlags;
+}

tensorboard/components/tf_feature_flags/tf-feature-flags-polymer.ts

@@ -0,0 +1,27 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+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 {customElement} from '@polymer/decorators';
+import {PolymerElement} from '@polymer/polymer';
+import * as tf_feature_flags from './feature-flags';
+
+/**
+ * This Polymer component allows the underlying API to be accessible from
+ * Angular TensorBoard by exposing otherwise mangled smybols.
+ */
+@customElement('tf-feature-flags')
+class TfFeatureFlags extends PolymerElement {
+  override _template = null;
+  tf_feature_flags = tf_feature_flags;
+}

tensorboard/webapp/BUILD

@@ -79,6 +79,9 @@ tf_ts_library(
     srcs = [
         "tb_polymer_interop_types.ts",
     ],
+    deps = [
+        "//tensorboard/webapp/feature_flag:types",
+    ],
 )
 
 tf_sass_binary(

tensorboard/webapp/feature_flag/effects/BUILD

@@ -8,6 +8,7 @@ tf_ng_module(
         "feature_flag_effects.ts",
     ],
     deps = [
+        "//tensorboard/webapp:tb_polymer_interop_types",
         "//tensorboard/webapp/feature_flag:force_svg_data_source",
         "//tensorboard/webapp/feature_flag/actions",
         "//tensorboard/webapp/feature_flag/store",

tensorboard/webapp/feature_flag/effects/feature_flag_effects.ts

@@ -16,17 +16,27 @@ limitations under the License.
 import {Injectable} from '@angular/core';
 import {Actions, createEffect, ofType} from '@ngrx/effects';
 import {Action, createAction, Store} from '@ngrx/store';
-import {combineLatestWith, map} from 'rxjs/operators';
+import {combineLatestWith, map, tap, withLatestFrom} from 'rxjs/operators';
+import '../../tb_polymer_interop_types';
 import {TBFeatureFlagDataSource} from '../../webapp_data_source/tb_feature_flag_data_source_types';
 import {partialFeatureFlagsLoaded} from '../actions/feature_flag_actions';
 import {ForceSvgDataSource} from '../force_svg_data_source';
-import {getIsAutoDarkModeAllowed} from '../store/feature_flag_selectors';
+import {
+  getFeatureFlags,
+  getIsAutoDarkModeAllowed,
+} from '../store/feature_flag_selectors';
 import {State} from '../store/feature_flag_types';
 
 const effectsInitialized = createAction('[FEATURE FLAG] Effects Init');
 
 @Injectable()
 export class FeatureFlagEffects {
+  // Ngrx assumes all Effect classes have properties that inherit from the base
+  // JS Object. `tf_feature_flags` does not, so we wrap it.
+  private readonly tfFeatureFlags = {
+    ref: document.createElement('tf-feature-flags').tf_feature_flags,
+  };
+
   /** @export */
   readonly getFeatureFlags$ = createEffect(() =>
     this.actions$.pipe(
@@ -46,6 +56,27 @@ export class FeatureFlagEffects {
     )
   );
 
+  /**
+   * Pass FeatureFlags to the Polymer portion of the code base immediately after
+   * feature flags have been finalized.
+   *
+   * @export
+   */
+  readonly updatePolymerFeatureFlags$ = createEffect(
+    () =>
+      this.actions$.pipe(
+        // partialFeatureFlagsLoaded triggers this effect but the actual
+        // feature flag values used are from the Store, given that it contains
+        // the finalized merged feature flags.
+        ofType(partialFeatureFlagsLoaded),
+        withLatestFrom(this.store.select(getFeatureFlags)),
+        tap(([, featureFlags]) => {
+          this.tfFeatureFlags.ref.setFeatureFlags(featureFlags);
+        })
+      ),
+    {dispatch: false}
+  );
+
   constructor(
     private readonly actions$: Actions,
     private readonly store: Store<State>,

tensorboard/webapp/feature_flag/effects/feature_flag_effects_test.ts

@@ -25,7 +25,10 @@ import {
 import {partialFeatureFlagsLoaded} from '../actions/feature_flag_actions';
 import {ForceSvgDataSource} from '../force_svg_data_source';
 import {ForceSvgDataSourceModule} from '../force_svg_data_source_module';
-import {getIsAutoDarkModeAllowed} from '../store/feature_flag_selectors';
+import {
+  getFeatureFlags,
+  getIsAutoDarkModeAllowed,
+} from '../store/feature_flag_selectors';
 import {State} from '../store/feature_flag_types';
 import {buildFeatureFlag} from '../testing';
 import {FeatureFlags} from '../types';
@@ -37,6 +40,7 @@ describe('feature_flag_effects', () => {
   let dataSource: TestingTBFeatureFlagDataSource;
   let forceSvgDataSource: ForceSvgDataSource;
   let effects: FeatureFlagEffects;
+  let setPolymerFeatureFlagsSpy: jasmine.Spy;
 
   beforeEach(async () => {
     actions = new ReplaySubject<Action>(1);
@@ -48,6 +52,16 @@ describe('feature_flag_effects', () => {
         provideMockStore(),
       ],
     }).compileComponents();
+
+    setPolymerFeatureFlagsSpy = jasmine.createSpy('setFeatureFlags');
+    const createElementSpy = spyOn(document, 'createElement');
+    createElementSpy.withArgs('tf-feature-flags').and.returnValue({
+      tf_feature_flags: {
+        setFeatureFlags: setPolymerFeatureFlagsSpy,
+      },
+    } as unknown as HTMLElement);
+    createElementSpy.and.callThrough();
+
     effects = TestBed.inject(FeatureFlagEffects);
     store = TestBed.inject<Store<State>>(Store) as MockStore<State>;
     dataSource = TestBed.inject(TestingTBFeatureFlagDataSource);
@@ -135,4 +149,32 @@ describe('feature_flag_effects', () => {
       expect(updateSpy).toHaveBeenCalledTimes(0);
     });
   });
+
+  describe('updatePolymerFeatureFlags$', () => {
+    it('sets polymer feature flags after data source fetch', () => {
+      // This represents the complete FeatureFlags object as calculated by the
+      // Store after the feature flags data source fetch.
+      store.overrideSelector(
+        getFeatureFlags,
+        buildFeatureFlag({inColab: true})
+      );
+      store.refreshState();
+
+      effects.updatePolymerFeatureFlags$.subscribe();
+
+      actions.next(
+        partialFeatureFlagsLoaded({
+          // This represents the incomplete FeatureFlags object that has just
+          // been fetched by the feature flags data source.
+          features: buildFeatureFlag({inColab: false}),
+        })
+      );
+
+      // Expect tf_feature_flags.setFeatureFlags() to be called using the
+      // FeatureFlags object from the Store (and not from the action).
+      expect(setPolymerFeatureFlagsSpy).toHaveBeenCalledOnceWith(
+        buildFeatureFlag({inColab: true})
+      );
+    });
+  });
 });

tensorboard/webapp/feature_flag/http/BUILD

@@ -2,12 +2,20 @@ load("//tensorboard/defs:defs.bzl", "tf_ng_module", "tf_ts_library")
 
 package(default_visibility = ["//tensorboard:internal"])
 
+tf_ng_module(
+    name = "const",
+    srcs = [
+        "const.ts",
+    ],
+)
+
 tf_ng_module(
     name = "http",
     srcs = [
         "feature_flag_http_interceptor.ts",
     ],
     deps = [
+        ":const",
         "//tensorboard/webapp/angular:expect_angular_common_http",
         "//tensorboard/webapp/feature_flag/store",
         "//tensorboard/webapp/feature_flag/store:types",
@@ -24,6 +32,7 @@ tf_ts_library(
         "feature_flag_http_interceptor_test.ts",
     ],
     deps = [
+        ":const",
         ":http",
         "//tensorboard/webapp/angular:expect_angular_common_http",
         "//tensorboard/webapp/angular:expect_angular_common_http_testing",

tensorboard/webapp/feature_flag/http/const.ts

@@ -0,0 +1,16 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+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.
+==============================================================================*/
+
+export const FEATURE_FLAGS_HEADER_NAME = 'X-TensorBoard-Feature-Flags';