feat: allow multiple subscriptions to single topic + auth

Signed-off-by: Heikki Hellgren <heikki.hellgren@op.fi>

plugins/signals-node/src/SignalsService.ts

@@ -28,6 +28,7 @@ import { JsonObject } from '@backstage/types';
 import {
   BackstageIdentityResponse,
   IdentityApi,
+  IdentityApiGetIdentityRequest,
 } from '@backstage/plugin-auth-node';
 
 /** @public */
@@ -160,6 +161,23 @@ export class SignalsService implements EventSubscriber {
       );
       connection.subscriptions.delete(message.topic as string);
     }
+
+    if (message.action === 'authenticate' && message.token) {
+      this.logger.info(`Connection ${connection.id} authenticated`);
+      this.identity
+        .getIdentity({
+          request: {
+            headers: { authorization: message.token },
+          },
+        } as IdentityApiGetIdentityRequest)
+        .then(identity => {
+          if (identity) {
+            connection.user = identity.identity.userEntityRef;
+            connection.ownershipEntityRefs =
+              identity.identity.ownershipEntityRefs;
+          }
+        });
+    }
   }
 
   /**
@@ -183,6 +201,11 @@ export class SignalsService implements EventSubscriber {
     message: JsonObject,
     brokedEvent: boolean,
   ) {
+    const jsonMessage = JSON.stringify({ topic, message });
+    if (jsonMessage.length === 0) {
+      return;
+    }
+
     this.connections.forEach(conn => {
       if (!conn.subscriptions.has(topic)) {
         return;
@@ -194,7 +217,7 @@ export class SignalsService implements EventSubscriber {
       ) {
         return;
       }
-      conn.ws.send(JSON.stringify({ topic, message }));
+      conn.ws.send(jsonMessage);
     });
 
     // If this event has not been broadcasted to all servers, then use

plugins/signals-react/api-report.md

@@ -11,8 +11,8 @@ export type SignalsApi = {
   subscribe(
     onMessage: (message: JsonObject, topic?: string) => void,
     topic: string,
-  ): void;
-  unsubscribe(topic: string): void;
+  ): string;
+  unsubscribe(subscription: string): void;
 };
 
 // @public (undocumented)

plugins/signals-react/src/api/SignalsApi.ts

@@ -26,7 +26,7 @@ export type SignalsApi = {
   subscribe(
     onMessage: (message: JsonObject, topic?: string) => void,
     topic: string,
-  ): void;
+  ): string;
 
-  unsubscribe(topic: string): void;
+  unsubscribe(subscription: string): void;
 };

plugins/signals-react/src/hooks/useSignalsApi.ts

@@ -16,7 +16,7 @@
 import { signalsApiRef } from '../api';
 import { useApi } from '@backstage/core-plugin-api';
 import { JsonObject } from '@backstage/types';
-import { useEffect } from 'react';
+import { useEffect, useState } from 'react';
 
 /** @public */
 export const useSignalsApi = (
@@ -24,13 +24,19 @@ export const useSignalsApi = (
   onMessage: (message: JsonObject) => void,
 ) => {
   const signals = useApi(signalsApiRef);
+  const [subscription, setSubscription] = useState<null | string>(null);
   useEffect(() => {
-    signals.subscribe(onMessage, topic);
-  }, [signals, onMessage, topic]);
+    if (!subscription) {
+      const sub = signals.subscribe(onMessage, topic);
+      setSubscription(sub);
+    }
+  }, [subscription, signals, onMessage, topic]);
 
   useEffect(() => {
     return () => {
-      signals.unsubscribe(topic);
+      if (subscription) {
+        signals.unsubscribe(subscription);
+      }
     };
-  }, [signals, topic]);
+  }, [subscription, signals, topic]);
 };

plugins/signals/api-report.md

@@ -5,19 +5,25 @@
 ```ts
 import { BackstagePlugin } from '@backstage/core-plugin-api';
 import { DiscoveryApi } from '@backstage/core-plugin-api';
+import { IdentityApi } from '@backstage/core-plugin-api';
 import { JsonObject } from '@backstage/types';
 import { SignalsApi } from '@backstage/plugin-signals-react';
 
 // @public (undocumented)
 export class SignalsClient implements SignalsApi {
   // (undocumented)
-  static create(options: { discoveryApi: DiscoveryApi }): SignalsClient;
+  static readonly CONNECT_TIMEOUT_MS: number;
   // (undocumented)
-  static instance: SignalsClient | null;
+  static create(options: {
+    identity: IdentityApi;
+    discoveryApi: DiscoveryApi;
+  }): SignalsClient;
   // (undocumented)
-  subscribe(onMessage: (message: JsonObject) => void, topic: string): void;
+  static readonly RECONNECT_TIMEOUT_MS: number;
   // (undocumented)
-  unsubscribe(topic: string): void;
+  subscribe(onMessage: (message: JsonObject) => void, topic: string): string;
+  // (undocumented)
+  unsubscribe(subscription: string): void;
 }
 
 // @public (undocumented)

plugins/signals/package.json

@@ -31,7 +31,8 @@
     "@material-ui/core": "^4.9.13",
     "@material-ui/icons": "^4.9.1",
     "@material-ui/lab": "^4.0.0-alpha.61",
-    "react-use": "^17.2.4"
+    "react-use": "^17.2.4",
+    "uuid": "^8.0.0"
   },
   "peerDependencies": {
     "react": "^16.13.1 || ^17.0.0"

plugins/signals/src/api/SignalsClient.ts

@@ -15,60 +15,105 @@
  */
 import { SignalsApi } from '@backstage/plugin-signals-react';
 import { JsonObject } from '@backstage/types';
-import { DiscoveryApi } from '@backstage/core-plugin-api';
+import { DiscoveryApi, IdentityApi } from '@backstage/core-plugin-api';
+import { v4 as uuid } from 'uuid';
+
+/** @internal */
+type Subscription = {
+  topic: string;
+  callback: (message: JsonObject) => void;
+};
+
+/** @internal */
+const WS_CLOSE_NORMAL = 1000;
+/** @internal */
+const WS_CLOSE_GOING_AWAY = 1001;
 
 /** @public */
 export class SignalsClient implements SignalsApi {
-  static instance: SignalsClient | null = null;
+  static readonly CONNECT_TIMEOUT_MS: number = 1000;
+  static readonly RECONNECT_TIMEOUT_MS: number = 5000;
   private ws: WebSocket | null = null;
-  private discoveryApi: DiscoveryApi;
-  private cbs: Map<string, (message: JsonObject) => void> = new Map();
-  private queue: JsonObject[] = [];
+  private subscriptions: Map<string, Subscription> = new Map();
+  private messageQueue: string[] = [];
   private reconnectTimeout: any;
 
-  static create(options: { discoveryApi: DiscoveryApi }) {
-    if (!SignalsClient.instance) {
-      SignalsClient.instance = new SignalsClient(options);
-    }
-    return SignalsClient.instance;
+  static create(options: {
+    identity: IdentityApi;
+    discoveryApi: DiscoveryApi;
+  }) {
+    const { identity, discoveryApi } = options;
+    return new SignalsClient(identity, discoveryApi);
   }
 
-  private constructor(options: { discoveryApi: DiscoveryApi }) {
-    this.discoveryApi = options.discoveryApi;
+  private constructor(
+    private identity: IdentityApi,
+    private discoveryApi: DiscoveryApi,
+  ) {}
+
+  subscribe(onMessage: (message: JsonObject) => void, topic: string): string {
+    const subscriptionId = uuid();
+    const exists = [...this.subscriptions.values()].find(
+      sub => sub.topic === topic,
+    );
+    this.subscriptions.set(subscriptionId, { topic, callback: onMessage });
+
+    this.connect()
+      .then(() => {
+        // Do not subscribe twice to same topic even there is multiple callbacks
+        if (!exists) {
+          this.send({ action: 'subscribe', topic });
+        }
+      })
+      .catch(() => {
+        this.reconnect();
+      });
+    return subscriptionId;
   }
 
-  subscribe(onMessage: (message: JsonObject) => void, topic: string): void {
-    // Do not allow to subscribe to same topic multiple times
-    if (this.cbs.has(topic)) {
+  unsubscribe(subscription: string): void {
+    const sub = this.subscriptions.get(subscription);
+    if (!sub) {
       return;
     }
+    const topic = sub.topic;
+    this.subscriptions.delete(subscription);
+    const exists = [...this.subscriptions.values()].find(
+      s => s.topic === topic,
+    );
+    // If there are subscriptions still listening to this topic, do not
+    // unsubscribe from the server
+    if (!exists) {
+      this.send({ action: 'unsubscribe', topic: sub.topic });
+    }
 
-    this.cbs.set(topic, onMessage);
-    this.connect().then(() => {
-      this.send({ action: 'subscribe', topic });
-    });
-  }
-
-  unsubscribe(topic: string): void {
-    this.cbs.delete(topic);
-    this.send({ action: 'unsubscribe', topic });
+    // If there are no subscriptions, close the connection
+    if (this.subscriptions.size === 0) {
+      this.ws?.close(WS_CLOSE_NORMAL);
+      this.ws = null;
+    }
   }
 
   private send(data?: JsonObject): void {
+    const jsonMessage = JSON.stringify(data);
+    if (jsonMessage.length === 0) {
+      return;
+    }
+
     if (!this.ws || this.ws.readyState !== WebSocket.OPEN) {
       if (data) {
-        this.queue.push(data);
+        this.messageQueue.unshift(jsonMessage);
       }
       return;
     }
 
     // First send queue
-    for (const msg of this.queue) {
-      this.ws!.send(JSON.stringify(msg));
+    for (const msg of this.messageQueue) {
+      this.ws!.send(msg);
     }
-    this.queue = [];
+    this.messageQueue = [];
     if (data) {
-      this.ws!.send(JSON.stringify(data));
+      this.ws!.send(jsonMessage);
     }
   }
 
@@ -83,31 +128,60 @@ export class SignalsClient implements SignalsApi {
     this.ws = new WebSocket(url.toString());
 
     this.ws.onmessage = (data: MessageEvent) => {
-      try {
-        const json = JSON.parse(data.data) as JsonObject;
-        if (json.topic) {
-          const cb = this.cbs.get(json.topic as string);
-          if (cb) {
-            cb(json.message as JsonObject);
-          }
-        }
-      } catch (e) {
-        // NOOP
-      }
+      this.handleMessage(data);
     };
 
     this.ws.onerror = () => {
       this.reconnect();
     };
 
-    this.ws.onclose = () => {
-      this.reconnect();
+    this.ws.onclose = (ev: CloseEvent) => {
+      if (ev.code !== WS_CLOSE_NORMAL && ev.code !== WS_CLOSE_GOING_AWAY) {
+        this.reconnect();
+      }
     };
 
-    while (this.ws.readyState !== WebSocket.OPEN) {
-      await new Promise(r => setTimeout(r, 10));
+    // Wait until connection is open
+    let connectSleep = 0;
+    while (
+      this.ws &&
+      this.ws.readyState !== WebSocket.OPEN &&
+      connectSleep < SignalsClient.CONNECT_TIMEOUT_MS
+    ) {
+      await new Promise(r => setTimeout(r, 100));
+      connectSleep += 100;
+    }
+
+    if (!this.ws || this.ws.readyState !== WebSocket.OPEN) {
+      throw new Error('Connect timeout');
+    }
+
+    // Authenticate
+    await this.authenticate();
+  }
+
+  private handleMessage(data: MessageEvent) {
+    try {
+      const json = JSON.parse(data.data) as JsonObject;
+      if (json.topic) {
+        for (const sub of this.subscriptions.values()) {
+          if (sub.topic === json.topic) {
+            sub.callback(json.message as JsonObject);
+          }
+        }
+      }
+    } catch (e) {
+      // NOOP
+    }
+  }
+
+  private async authenticate() {
+    const { token } = await this.identity.getCredentials();
+    if (token) {
+      // Authentication is done with websocket message to server as the plain
+      // websocket does not allow sending headers during connection upgrade
+      this.send({ action: 'authenticate', token: token });
     }
-    this.send();
   }
 
   private reconnect() {
@@ -116,16 +190,21 @@ export class SignalsClient implements SignalsApi {
     }
 
     this.reconnectTimeout = setTimeout(() => {
+      this.reconnectTimeout = null;
       if (this.ws) {
         this.ws.close();
       }
       this.ws = null;
-      this.connect().then(() => {
-        // Resubscribe to existing topics in case we lost connection
-        for (const topic of this.cbs.keys()) {
-          this.send({ action: 'subscribe', topic });
-        }
-      });
-    }, 5000);
+      this.connect()
+        .then(() => {
+          // Resubscribe to existing topics in case we lost connection
+          for (const topic of this.subscriptions.keys()) {
+            this.send({ action: 'subscribe', topic });
+          }
+        })
+        .catch(() => {
+          this.reconnect();
+        });
+    }, SignalsClient.RECONNECT_TIMEOUT_MS);
   }
 }

plugins/signals/src/plugin.ts

@@ -17,6 +17,7 @@ import {
   createApiFactory,
   createPlugin,
   discoveryApiRef,
+  identityApiRef,
 } from '@backstage/core-plugin-api';
 import { signalsApiRef } from '@backstage/plugin-signals-react';
 import { SignalsClient } from './api/SignalsClient';
@@ -28,10 +29,12 @@ export const signalsPlugin = createPlugin({
     createApiFactory({
       api: signalsApiRef,
       deps: {
+        identity: identityApiRef,
         discoveryApi: discoveryApiRef,
       },
-      factory: ({ discoveryApi }) =>
+      factory: ({ identity, discoveryApi }) =>
         SignalsClient.create({
+          identity,
           discoveryApi,
         }),
     }),

yarn.lock

@@ -8914,6 +8914,7 @@ __metadata:
     "@testing-library/user-event": ^14.0.0
     msw: ^1.0.0
     react-use: ^17.2.4
+    uuid: ^8.0.0
   peerDependencies:
     react: ^16.13.1 || ^17.0.0
   languageName: unknown