Added trace plugin framework.

Change-Id: Icf9286fe6b64c9714427b2d37f41a81708d97d5c

12 files changed, 282 insertions, 11 deletions

diff --git a/tools/gen_ui_imports b/tools/gen_ui_imports
index b644dfea5..36b3825cd 100755
--- a/tools/gen_ui_imports
+++ b/tools/gen_ui_imports
@@ -36,8 +36,7 @@ from __future__ import print_function
 
 import os
 import argparse
-import subprocess
-import sys
+import re
 
 ROOT_DIR = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
 UI_SRC_DIR = os.path.join(ROOT_DIR, 'ui', 'src')
@@ -45,7 +44,8 @@ PLUGINS_PATH = os.path.join(UI_SRC_DIR, 'common', 'plugins')
 
 
 def to_camel_case(s):
-  first, *rest = s.split('_')
+  # Split string on periods and underscores
+  first, *rest = re.split(r'\.|\_', s)
   return first + ''.join(x.title() for x in rest)
 
 
@@ -70,7 +70,7 @@ def gen_imports(input_dir, output_path):
   import_text = '\n'.join(imports)
   registration_text = '\n'.join(registrations)
 
-  expected = f"{header}\n\n{import_text}\n\n{registration_text}"
+  expected = f"{header}\n\n{import_text}\n\n{registration_text}\n"
 
   with open(output_path, 'w') as f:
     f.write(expected)
diff --git a/ui/build.js b/ui/build.js
index 92f793bcb..f8e2d903a 100644
--- a/ui/build.js
+++ b/ui/build.js
@@ -237,6 +237,7 @@ async function main() {
     scanDir('buildtools/typefaces');
     scanDir('buildtools/catapult_trace_viewer');
     generateImports('ui/src/tracks', 'all_tracks.ts');
+    generateImports('ui/src/plugins', 'all_plugins.ts');
     compileProtos();
     genVersion();
     transpileTsProject('ui');
diff --git a/ui/src/common/empty_state.ts b/ui/src/common/empty_state.ts
index af1785f52..6044708d1 100644
--- a/ui/src/common/empty_state.ts
+++ b/ui/src/common/empty_state.ts
@@ -173,5 +173,8 @@ export function createEmptyState(): State {
       textEntry: '',
       hideNonMatching: true,
     },
+
+    // Somewhere to store plugins' persistent state.
+    plugins: {},
   };
 }
diff --git a/ui/src/common/plugin_api.ts b/ui/src/common/plugin_api.ts
index 9ba3a9932..53e09a2e1 100644
--- a/ui/src/common/plugin_api.ts
+++ b/ui/src/common/plugin_api.ts
@@ -16,6 +16,8 @@ import {EngineProxy} from '../common/engine';
 import {TrackControllerFactory} from '../controller/track_controller';
 import {TrackCreator} from '../frontend/track';
 
+import {TracePluginFactory} from './plugins';
+
 export {EngineProxy} from '../common/engine';
 export {
   LONG,
@@ -68,6 +70,10 @@ export interface PluginContext {
   // could be registered in dev.perfetto.CounterTrack - a whole
   // different plugin.
   registerTrack(track: TrackCreator): void;
+
+  // Register a new plugin factory for a plugin whose lifecycle in linked to
+  // that of the trace.
+  registerTracePluginFactory<T>(pluginFactory: TracePluginFactory<T>): void;
 }
 
 export interface PluginInfo {
diff --git a/ui/src/common/plugins.ts b/ui/src/common/plugins.ts
index aa414bc74..5a362c792 100644
--- a/ui/src/common/plugins.ts
+++ b/ui/src/common/plugins.ts
@@ -12,21 +12,57 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import {Engine} from '../common/engine';
+import {Disposable} from 'src/base/disposable';
+
 import {
   TrackControllerFactory,
   trackControllerRegistry,
 } from '../controller/track_controller';
+import {Store} from '../frontend/store';
 import {TrackCreator} from '../frontend/track';
 import {trackRegistry} from '../frontend/track_registry';
 
+import {Engine} from './engine';
 import {
+  EngineProxy,
   PluginContext,
   PluginInfo,
   TrackInfo,
   TrackProvider,
 } from './plugin_api';
 import {Registry} from './registry';
+import {State} from './state';
+
+// All trace plugins must implement this interface.
+export interface TracePlugin extends Disposable {
+  // This is where we would add potential extension points that plugins can
+  // override.
+  // E.g. commands(): Command[];
+  // For now, plugins don't do anything so this interface is empty.
+}
+
+// This interface defines what a plugin factory should look like.
+// This can be defined in the plugin class definition by defining a constructor
+// and the relevant static methods:
+// E.g.
+// class MyPlugin implements TracePlugin<MyState> {
+//   static migrate(initialState: unknown): MyState {...}
+//   constructor(store: Store<MyState>, engine: EngineProxy) {...}
+//   ... methods from the TracePlugin interface go here ...
+// }
+// ... which can then be passed around by class i.e. MyPlugin
+export interface TracePluginFactory<StateT> {
+  // Function to migrate the persistent state. Called before new().
+  migrate(initialState: unknown): StateT;
+
+  // Instantiate the plugin.
+  new(store: Store<StateT>, engine: EngineProxy): TracePlugin;
+}
+
+interface TracePluginContext {
+  plugin: TracePlugin;
+  store: Store<unknown>;
+}
 
 // Every plugin gets its own PluginContext. This is how we keep track
 // what each plugin is doing and how we can blame issues on particular
@@ -34,6 +70,8 @@ import {Registry} from './registry';
 export class PluginContextImpl implements PluginContext {
   readonly pluginId: string;
   private trackProviders: TrackProvider[];
+  private tracePluginFactory?: TracePluginFactory<any>;
+  private _tracePluginCtx?: TracePluginContext;
 
   constructor(pluginId: string) {
     this.pluginId = pluginId;
@@ -53,6 +91,10 @@ export class PluginContextImpl implements PluginContext {
   registerTrackProvider(provider: TrackProvider) {
     this.trackProviders.push(provider);
   }
+
+  registerTracePluginFactory<T>(pluginFactory: TracePluginFactory<T>): void {
+    this.tracePluginFactory = pluginFactory;
+  }
   // ==================================================================
 
   // ==================================================================
@@ -62,11 +104,50 @@ export class PluginContextImpl implements PluginContext {
     return this.trackProviders.map((f) => f(proxy));
   }
 
+  onTraceLoad(store: Store<State>, engine: Engine): void {
+    const TracePluginClass = this.tracePluginFactory;
+    if (TracePluginClass) {
+      // Make an engine proxy for this plugin.
+      const engineProxy = engine.getProxy(this.pluginId);
+
+      // Extract the initial state and pass to the plugin factory for migration.
+      const initialState = store.state.plugins[this.pluginId];
+      const migratedState = TracePluginClass.migrate(initialState);
+
+      // Store the initial state in our root store.
+      store.edit((draft) => {
+        draft.plugins[this.pluginId] = migratedState;
+      });
+
+      // Create a proxy store for our plugin to use.
+      const storeProxy = store.createProxy<unknown>(['plugins', this.pluginId]);
+
+      // Instantiate the plugin.
+      this._tracePluginCtx = {
+        plugin: new TracePluginClass(storeProxy, engineProxy),
+        store: storeProxy,
+      };
+    }
+  }
+
+  onTraceClosed() {
+    if (this._tracePluginCtx) {
+      this._tracePluginCtx.plugin.dispose();
+      this._tracePluginCtx.store.dispose();
+      this._tracePluginCtx = undefined;
+    }
+  }
+
+  get tracePlugin(): TracePlugin|undefined {
+    return this._tracePluginCtx?.plugin;
+  }
+
   // Unload the plugin. Ideally no plugin code runs after this point.
   // PluginContext should unregister everything.
   revoke() {
     // TODO(hjd): Remove from trackControllerRegistry, trackRegistry,
     // etc.
+    // TODO(stevegolton): Close the trace plugin.
   }
   // ==================================================================
 }
@@ -123,6 +204,18 @@ export class PluginManager {
     }
     return promises;
   }
+
+  onTraceLoad(store: Store<State>, engine: Engine): void {
+    for (const context of this.contexts.values()) {
+      context.onTraceLoad(store, engine);
+    }
+  }
+
+  onTraceClose() {
+    for (const context of this.contexts.values()) {
+      context.onTraceClosed();
+    }
+  }
 }
 
 // TODO(hjd): Sort out the story for global singletons like these:
diff --git a/ui/src/common/state.ts b/ui/src/common/state.ts
index dba115801..5f5e1a42f 100644
--- a/ui/src/common/state.ts
+++ b/ui/src/common/state.ts
@@ -624,6 +624,9 @@ export interface State {
   // Pending deeplink which will happen when we first finish opening a
   // trace.
   pendingDeeplink?: PendingDeeplinkState;
+
+  // Individual plugin states
+  plugins: {[key: string]: any};
 }
 
 export const defaultTraceTime = {
diff --git a/ui/src/controller/trace_controller.ts b/ui/src/controller/trace_controller.ts
index 047f00558..d46a6c0ae 100644
--- a/ui/src/controller/trace_controller.ts
+++ b/ui/src/controller/trace_controller.ts
@@ -30,6 +30,7 @@ import {
   getEnabledMetatracingCategories,
   isMetatracingEnabled,
 } from '../common/metatracing';
+import {pluginManager} from '../common/plugins';
 import {
   LONG,
   NUM,
@@ -45,11 +46,7 @@ import {
   PendingDeeplinkState,
   ProfileType,
 } from '../common/state';
-import {Span} from '../common/time';
-import {
-  TPTime,
-  TPTimeSpan,
-} from '../common/time';
+import {Span, TPTime, TPTimeSpan} from '../common/time';
 import {resetEngineWorker, WasmEngineProxy} from '../common/wasm_engine_proxy';
 import {BottomTabList} from '../frontend/bottom_tab';
 import {
@@ -342,6 +339,7 @@ export class TraceController extends Controller<States> {
   }
 
   onDestroy() {
+    pluginManager.onTraceClose();
     globals.engines.delete(this.engineId);
   }
 
@@ -556,6 +554,8 @@ export class TraceController extends Controller<States> {
       }
     }
 
+    pluginManager.onTraceLoad(globals.store, engine);
+
     return engineMode;
   }
 
diff --git a/ui/src/frontend/debug.ts b/ui/src/frontend/debug.ts
index 7e9c9e523..64de18a19 100644
--- a/ui/src/frontend/debug.ts
+++ b/ui/src/frontend/debug.ts
@@ -16,6 +16,7 @@ import {produce} from 'immer';
 import m from 'mithril';
 
 import {Actions} from '../common/actions';
+import {pluginManager} from '../common/plugins';
 import {getSchema} from '../common/schema';
 
 import {globals} from './globals';
@@ -28,6 +29,7 @@ declare global {
     globals: typeof globals;
     Actions: typeof Actions;
     produce: typeof produce;
+    pluginManager: typeof pluginManager
   }
 }
 
@@ -37,4 +39,5 @@ export function registerDebugGlobals() {
   window.globals = globals;
   window.Actions = Actions;
   window.produce = produce;
+  window.pluginManager = pluginManager;
 }
diff --git a/ui/src/frontend/index.ts b/ui/src/frontend/index.ts
index 9a8d3a488..22a512290 100644
--- a/ui/src/frontend/index.ts
+++ b/ui/src/frontend/index.ts
@@ -14,6 +14,7 @@
 
 // Keep this import first.
 import '../core/static_initializers';
+import '../gen/all_plugins';
 
 import {Draft} from 'immer';
 import m from 'mithril';
diff --git a/ui/src/frontend/store.ts b/ui/src/frontend/store.ts
index a178988b6..31fd046da 100644
--- a/ui/src/frontend/store.ts
+++ b/ui/src/frontend/store.ts
@@ -140,16 +140,19 @@ class StoreImpl<T> implements Store<T> {
 export class ProxyStoreImpl<RootT, T> implements Store<T> {
   private subscriptions = new Set<SubscriptionCallback<T>>();
   private rootSubscription;
+  private rootStore?: Store<RootT>;
 
   constructor(
-      private rootStore: Store<RootT>,
+      rootStore: Store<RootT>,
       private path: Path,
   ) {
+    this.rootStore = rootStore;
     this.rootSubscription = rootStore.subscribe(this.rootUpdateHandler);
   }
 
   dispose() {
     this.rootSubscription.dispose();
+    this.rootStore = undefined;
   }
 
   private rootUpdateHandler = (newState: RootT, oldState: RootT) => {
@@ -164,10 +167,15 @@ export class ProxyStoreImpl<RootT, T> implements Store<T> {
   };
 
   get state(): T {
+    if (!this.rootStore) {
+      throw new StoreError('Proxy store is no longer useable');
+    }
+
     const state = lookupPath<T, RootT>(this.rootStore.state, this.path);
     if (state === undefined) {
       throw new StoreError(`No such subtree: ${this.path}`);
     }
+
     return state;
   }
 
@@ -180,6 +188,10 @@ export class ProxyStoreImpl<RootT, T> implements Store<T> {
   }
 
   private applyEdits(edits: Edit<T>[]): void {
+    if (!this.rootStore) {
+      throw new StoreError('Proxy store is no longer useable');
+    }
+
     // Transform edits to work on the root store.
     const rootEdits = edits.map(
         (edit) => (state: Draft<RootT>) => {
@@ -198,6 +210,10 @@ export class ProxyStoreImpl<RootT, T> implements Store<T> {
   }
 
   createProxy<NewSubStateT>(path: Path): Store<NewSubStateT> {
+    if (!this.rootStore) {
+      throw new StoreError('Proxy store is no longer useable');
+    }
+
     const fullPath = [...this.path, ...path];
     return new ProxyStoreImpl<RootT, NewSubStateT>(this.rootStore, fullPath);
   }
diff --git a/ui/src/frontend/store_unittest.ts b/ui/src/frontend/store_unittest.ts
index 13395a1df..01d39d123 100644
--- a/ui/src/frontend/store_unittest.ts
+++ b/ui/src/frontend/store_unittest.ts
@@ -282,4 +282,104 @@ describe('proxy store', () => {
     // Ensure proxy callback hasn't been called
     expect(callback).not.toHaveBeenCalled();
   });
+
+  it('notifies subscribers', () => {
+    const store = createStore(initialState);
+    const fooState = store.createProxy<Foo>(['foo']);
+    const callback = jest.fn();
+
+    fooState.subscribe(callback);
+
+    store.edit((draft) => {
+      draft.foo.counter += 1;
+    });
+
+    expect(callback).toHaveBeenCalledTimes(1);
+    expect(callback).toHaveBeenCalledWith(
+        {
+          ...initialState.foo,
+          counter: 1,
+        },
+        initialState.foo);
+  });
+
+  it('does not notify unsubscribed subscribers', () => {
+    const store = createStore(initialState);
+    const fooState = store.createProxy<Foo>(['foo']);
+    const callback = jest.fn();
+
+    // Subscribe then immediately unsubscribe
+    fooState.subscribe(callback).dispose();
+
+    // Make an arbitrary edit
+    fooState.edit((draft) => {
+      draft.counter += 1;
+    });
+
+    expect(callback).not.toHaveBeenCalled();
+  });
+
+  it('throws on state access when path doesn\'t exist', () => {
+    const store = createStore(initialState);
+
+    // This path is incorrect - baz doesn't exist in State
+    const fooStore = store.createProxy<Foo>(['baz']);
+
+    expect(() => {
+      fooStore.state;
+    }).toThrow(StoreError);
+  });
+
+  it('throws on edit when path doesn\'t exist', () => {
+    const store = createStore(initialState);
+
+    // This path is incorrect - baz doesn't exist in State
+    const fooState = store.createProxy<Foo>(['baz']);
+
+    expect(() => {
+      fooState.edit((draft) => {
+        draft.counter += 1;
+      });
+    }).toThrow(StoreError);
+  });
+
+  it('notifies when relevant edits are made from root store', () => {
+    const store = createStore(initialState);
+    const fooState = store.createProxy<Foo>(['foo']);
+    const callback = jest.fn();
+
+    // Subscribe on the proxy store
+    fooState.subscribe(callback);
+
+    // Edit the root store
+    store.edit((draft) => {
+      draft.foo.counter++;
+    });
+
+    // Expect proxy callback called with correct subtree
+    expect(callback).toHaveBeenCalled();
+    expect(callback).toHaveBeenCalledWith(
+        {
+          ...initialState.foo,
+          counter: 1,
+        },
+        initialState.foo);
+  });
+
+  it('ignores irrrelevant edits from the root store', () => {
+    const store = createStore(initialState);
+    const nestedStore = store.createProxy<NestedState>(['foo', 'nested']);
+    const callback = jest.fn();
+
+    // Subscribe on the proxy store
+    nestedStore.subscribe(callback);
+
+    // Edit an irrelevant subtree on the root store
+    store.edit((draft) => {
+      draft.foo.counter++;
+    });
+
+    // Ensure proxy callback hasn't been called
+    expect(callback).not.toHaveBeenCalled();
+  });
 });
diff --git a/ui/src/plugins/dev.perfetto.ExamplePlugin/index.ts b/ui/src/plugins/dev.perfetto.ExamplePlugin/index.ts
new file mode 100644
index 000000000..3305a90f6
--- /dev/null
+++ b/ui/src/plugins/dev.perfetto.ExamplePlugin/index.ts
@@ -0,0 +1,45 @@
+// Copyright (C) 2023 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 {EngineProxy, PluginContext} from '../../common/plugin_api';
+import {TracePlugin} from '../../common/plugins';
+import {Store} from '../../frontend/store';
+
+interface ExampleState {
+  counter: number;
+}
+
+// This is just an example plugin, used to prove that the plugin system works.
+class ExamplePlugin implements TracePlugin {
+  static migrate(_initialState: unknown): ExampleState {
+    return {counter: 0};
+  }
+
+  constructor(_store: Store<ExampleState>, _engine: EngineProxy) {
+    // No-op
+  }
+
+  dispose(): void {
+    // No-op
+  }
+}
+
+function activate(ctx: PluginContext) {
+  ctx.registerTracePluginFactory(ExamplePlugin);
+}
+
+export const plugin = {
+  pluginId: 'dev.perfetto.ExamplePlugin',
+  activate,
+};