Implement receiving messages from add-on on the client

Reviewed By: mweststrate

Differential Revision: D34249101

fbshipit-source-id: 07297b84ed8640e3b41599726ba613b6b4e2b62e

desktop/flipper-common/src/ServerAddOn.tsx

@@ -21,6 +21,17 @@ export interface ServerAddOnControls {
     method: string,
     params?: unknown,
   ) => Promise<object | string | number | boolean | null>;
+  receiveMessage: (
+    pluginName: string,
+    method: string,
+    receiver: (data: unknown) => void,
+  ) => void;
+  receiveAnyMessage: (
+    pluginName: string,
+    receiver: (method: string, data: unknown) => void,
+  ) => void;
+  unsubscribePlugin: (pluginName: string) => void;
+  unsubscribe: () => void;
 }
 
 // TODO: Share with js-flipper? Is it worth it?

desktop/flipper-plugin/src/plugin/PluginBase.tsx

@@ -386,11 +386,18 @@ export abstract class BasePluginInstance {
           method as string,
           params,
         ),
-      onServerAddOnMessage: (_event, _cb) => {
-        // TODO: Implement me
+      onServerAddOnMessage: (event, cb) => {
+        this.serverAddOnControls.receiveMessage(
+          this.definition.packageName,
+          event as string,
+          batched(cb),
+        );
       },
-      onServerAddOnUnhandledMessage: (_cb) => {
-        // TODO: Implement me
+      onServerAddOnUnhandledMessage: (cb) => {
+        this.serverAddOnControls.receiveAnyMessage(
+          this.definition.packageName,
+          batched(cb),
+        );
       },
     };
   }
@@ -436,6 +443,7 @@ export abstract class BasePluginInstance {
     this.crashListeners.splice(0).forEach((handle) => {
       this.device.removeCrashListener(handle);
     });
+    this.serverAddOnControls.unsubscribePlugin(this.definition.packageName);
     this.events.emit('destroy');
     this.destroyed = true;
   }

desktop/flipper-plugin/src/test-utils/test-utils.tsx

@@ -648,5 +648,9 @@ function createServerAddOnControlsMock(): ServerAddOnControls {
     start: createStubFunction(),
     stop: createStubFunction(),
     sendMessage: createStubFunction(),
+    receiveMessage: createStubFunction(),
+    receiveAnyMessage: createStubFunction(),
+    unsubscribePlugin: createStubFunction(),
+    unsubscribe: createStubFunction(),
   };
 }

desktop/flipper-ui-core/src/Client.tsx

@@ -126,7 +126,7 @@ export default class Client extends EventEmitter {
   > = {};
   sandyPluginStates = new Map<string /*pluginID*/, _SandyPluginInstance>();
   private readonly serverAddOnControls: ServerAddOnControls;
-  private readonly flipperServer: Pick<FlipperServer, 'exec'>;
+  private readonly flipperServer: FlipperServer;
 
   constructor(
     id: string,
@@ -136,7 +136,7 @@ export default class Client extends EventEmitter {
     store: Store,
     plugins: Plugins | null | undefined,
     device: BaseDevice,
-    flipperServer: Pick<FlipperServer, 'exec'>,
+    flipperServer: FlipperServer,
   ) {
     super();
     this.connected.set(!!conn);
@@ -281,6 +281,7 @@ export default class Client extends EventEmitter {
   destroy() {
     this.disconnect();
     this.plugins.forEach((pluginId) => this.stopPluginIfNeeded(pluginId, true));
+    this.serverAddOnControls.unsubscribe();
   }
 
   // gets a plugin by pluginId

desktop/flipper-ui-core/src/__tests__/disconnect.node.tsx

@@ -226,13 +226,12 @@ test('new clients replace old ones', async () => {
 
   const client2 = await createClient(device, 'AnotherApp', client.query, true);
   await handleClientConnected(
-    {
-      exec: (async () => {
-        return {
-          success: {}, // {plugins: []},
-        };
-      }) as any,
-    },
+    TestUtils.createFlipperServerMock({
+      'client-request-response': async () => ({
+        success: [],
+        length: 0,
+      }),
+    }),
     store,
     logger,
     client2,

desktop/flipper-ui-core/src/devices/BaseDevice.tsx

@@ -390,5 +390,6 @@ export default class BaseDevice implements Device {
       instance.destroy();
     });
     this.sandyPluginStates.clear();
+    this.serverAddOnControls.unsubscribe();
   }
 }

desktop/flipper-ui-core/src/dispatcher/flipperServer.tsx

@@ -240,7 +240,7 @@ function handleDeviceConnected(
 }
 
 export async function handleClientConnected(
-  server: Pick<FlipperServer, 'exec'>,
+  server: FlipperServer,
   store: Store,
   logger: Logger,
   {id, query}: ClientDescription,

desktop/flipper-ui-core/src/utils/createServerAddOnControls.tsx

@@ -8,35 +8,87 @@
  */
 
 import {
+  ExecuteMessage,
   FlipperServer,
   ServerAddOnControls,
   deserializeRemoteError,
 } from 'flipper-common';
 
-export const createServerAddOnControls = (
-  flipperServer: Pick<FlipperServer, 'exec'>,
-): ServerAddOnControls => ({
-  start: (pluginName, owner) =>
-    flipperServer.exec('plugins-server-add-on-start', pluginName, owner),
-  stop: (pluginName, owner) =>
-    flipperServer.exec('plugins-server-add-on-stop', pluginName, owner),
-  sendMessage: async (pluginName, method, params) => {
-    const res = await flipperServer.exec(
-      'plugins-server-add-on-request-response',
-      {
-        method: 'execute',
-        params: {
-          method,
-          api: pluginName,
-          params,
-        },
-      },
-    );
+type PluginName = string;
+type Method = string;
 
-    if (res.error) {
-      throw deserializeRemoteError(res.error);
+export const createServerAddOnControls = (
+  flipperServer: FlipperServer,
+): ServerAddOnControls => {
+  const methodHandlers = new Map<
+    PluginName,
+    Map<Method, (data: unknown) => void>
+  >();
+  const catchAllHandlers = new Map<
+    PluginName,
+    (method: string, data: unknown) => void
+  >();
+
+  let subscribed = false;
+  const subscriptionCb = ({params}: ExecuteMessage) => {
+    const pluginName = params.api;
+
+    const methodHandler = methodHandlers.get(pluginName)?.get(params.method);
+
+    if (methodHandler) {
+      methodHandler(params.params);
+      return;
     }
 
-    return res.success;
-  },
-});
+    const catchAllHandler = catchAllHandlers.get(pluginName);
+    catchAllHandler?.(params.method, params.params);
+  };
+
+  return {
+    start: (pluginName, owner) =>
+      flipperServer.exec('plugins-server-add-on-start', pluginName, owner),
+    stop: (pluginName, owner) =>
+      flipperServer.exec('plugins-server-add-on-stop', pluginName, owner),
+    sendMessage: async (pluginName, method, params) => {
+      const res = await flipperServer.exec(
+        'plugins-server-add-on-request-response',
+        {
+          method: 'execute',
+          params: {
+            method,
+            api: pluginName,
+            params,
+          },
+        },
+      );
+
+      if (res.error) {
+        throw deserializeRemoteError(res.error);
+      }
+
+      return res.success;
+    },
+    receiveMessage: (pluginName, method, receiver) => {
+      if (!methodHandlers.has(pluginName)) {
+        methodHandlers.set(pluginName, new Map());
+      }
+      methodHandlers.get(pluginName)!.set(method, receiver);
+
+      // Subscribe client/device to messages from flipper server only when the first plugin subscribes to them
+      if (!subscribed) {
+        subscribed = true;
+        flipperServer.on('plugins-server-add-on-message', subscriptionCb);
+      }
+    },
+    receiveAnyMessage: (pluginName, receiver) => {
+      catchAllHandlers.set(pluginName, receiver);
+    },
+    unsubscribePlugin: (pluginName) => {
+      methodHandlers.delete(pluginName);
+      catchAllHandlers.delete(pluginName);
+    },
+    unsubscribe: () => {
+      flipperServer.off('plugins-server-add-on-message', subscriptionCb);
+    },
+  };
+};