feat: add onConnectionStateChange callback for flux subscription (#2782)

* Add onConnectionStateChange callback for flux subscription Fixes: #2769 * Add onConnectionStateChange callback for flux subscription Fixes: #2769 * Fix tests * Fix tests * Fix signals test Fixes: #2769 * Fix formatting Fixes: #2769 --------- Co-authored-by: Anton Platonov <[email protected]>
vaadin · Oct 16, 2024 · 9bafcb7 · 9bafcb7
1 parent 0daf4dc
commit 9bafcb7
Show file tree

Hide file tree

Showing 5 changed files with 131 additions and 5 deletions.
diff --git a/packages/ts/frontend/src/Connect.ts b/packages/ts/frontend/src/Connect.ts
@@ -9,7 +9,11 @@ import {
   UnauthorizedResponseError,
   type ValidationErrorData,
 } from './EndpointErrors.js';
-import { type ActionOnLostSubscription, FluxConnection } from './FluxConnection.js';
+import {
+  type ActionOnLostSubscription,
+  FluxConnection,
+  type FluxSubscriptionStateChangeEvent,
+} from './FluxConnection.js';
 import type { VaadinWindow } from './types.js';
 
 const $wnd = window as VaadinWindow;
@@ -43,6 +47,9 @@ export interface Subscription<T> {
   /** Called when a new value is available. */
   onNext(callback: (value: T) => void): Subscription<T>;
 
+  /** Called when the subscription state changes. */
+  onConnectionStateChange(callback: (event: FluxSubscriptionStateChangeEvent) => void): Subscription<T>;
+
   /**
    * Called when the connection is restored, but there's no longer a valid subscription. If the callback returns
    * `ActionOnLostSubscription.RESUBSCRIBE`, the subscription will be re-established by connecting to the same

diff --git a/packages/ts/frontend/src/FluxConnection.ts b/packages/ts/frontend/src/FluxConnection.ts
@@ -41,6 +41,29 @@ export enum ActionOnLostSubscription {
   REMOVE = 'remove',
 }
 
+/**
+ * Possible states of a flux subscription.
+ */
+export enum FluxSubscriptionState {
+  /**
+   * The subscription is not connected and is trying to connect.
+   */
+  CONNECTING = 'connecting',
+  /**
+   * The subscription is connected and receiving updates.
+   */
+  CONNECTED = 'connected',
+  /**
+   * The subscription is closed and is not trying to reconnect.
+   */
+  CLOSED = 'closed',
+}
+
+/**
+ * Event wrapper for flux subscription connection state change callback
+ */
+export type FluxSubscriptionStateChangeEvent = CustomEvent<{ state: FluxSubscriptionState }>;
+
 type EndpointInfo = {
   endpointName: string;
   methodName: string;
@@ -58,6 +81,8 @@ export class FluxConnection extends EventTarget {
   readonly #onCompleteCallbacks = new Map<string, () => void>();
   readonly #onErrorCallbacks = new Map<string, () => void>();
   readonly #onNextCallbacks = new Map<string, (value: any) => void>();
+  readonly #onStateChangeCallbacks = new Map<string, (event: FluxSubscriptionStateChangeEvent) => void>();
+  readonly #statusOfSubscriptions = new Map<string, FluxSubscriptionState>();
   #pendingMessages: ServerMessage[] = [];
   #socket?: Atmosphere.Request;
 
@@ -82,6 +107,7 @@ export class FluxConnection extends EventTarget {
     const msg: ServerConnectMessage = { '@type': 'subscribe', endpointName, id, methodName, params };
     this.#send(msg);
     this.#endpointInfos.set(id, { endpointName, methodName, params });
+    this.#setSubscriptionConnState(id, FluxSubscriptionState.CONNECTING);
     const hillaSubscription: Subscription<any> = {
       cancel: () => {
         if (!this.#endpointInfos.has(id)) {
@@ -121,6 +147,13 @@ export class FluxConnection extends EventTarget {
         }
         return hillaSubscription;
       },
+      onConnectionStateChange: (callback: (event: FluxSubscriptionStateChangeEvent) => void): Subscription<any> => {
+        this.#onStateChangeCallbacks.set(id, callback);
+        callback(
+          new CustomEvent('subscription-state-change', { detail: { state: this.#statusOfSubscriptions.get(id)! } }),
+        );
+        return hillaSubscription;
+      },
     };
     return hillaSubscription;
   }
@@ -193,18 +226,41 @@ export class FluxConnection extends EventTarget {
       onReconnect: () => {
         if (this.state !== State.RECONNECTING) {
           this.state = State.RECONNECTING;
+          this.#endpointInfos.forEach((endpointInfo, id) => {
+            if (endpointInfo.reconnect?.() === ActionOnLostSubscription.RESUBSCRIBE) {
+              this.#setSubscriptionConnState(id, FluxSubscriptionState.CONNECTING);
+            } else {
+              this.#setSubscriptionConnState(id, FluxSubscriptionState.CLOSED);
+            }
+          });
         }
       },
       onFailureToReconnect: () => {
         if (this.state !== State.INACTIVE) {
           this.state = State.INACTIVE;
           this.dispatchEvent(new CustomEvent('state-changed', { detail: { active: false } }));
+          this.#endpointInfos.forEach((_, id) => this.#setSubscriptionConnState(id, FluxSubscriptionState.CLOSED));
         }
       },
       ...atmosphereOptions,
     } satisfies Atmosphere.Request);
   }
 
+  #setSubscriptionConnState(id: string, state: FluxSubscriptionState) {
+    const currentState = this.#statusOfSubscriptions.get(id);
+    if (!currentState) {
+      this.#statusOfSubscriptions.set(id, state);
+      this.#onStateChangeCallbacks.get(id)?.(
+        new CustomEvent('subscription-state-change', { detail: { state: this.#statusOfSubscriptions.get(id)! } }),
+      );
+    } else if (currentState !== state) {
+      this.#statusOfSubscriptions.set(id, state);
+      this.#onStateChangeCallbacks.get(id)?.(
+        new CustomEvent('subscription-state-change', { detail: { state: this.#statusOfSubscriptions.get(id)! } }),
+      );
+    }
+  }
+
   #handleMessage(message: unknown) {
     if (isClientMessage(message)) {
       const { id } = message;
@@ -215,6 +271,7 @@ export class FluxConnection extends EventTarget {
         if (callback) {
           callback(message.item);
         }
+        this.#setSubscriptionConnState(id, FluxSubscriptionState.CONNECTED);
       } else if (message['@type'] === 'complete') {
         this.#onCompleteCallbacks.get(id)?.();
         this.#removeSubscription(id);
@@ -238,6 +295,9 @@ export class FluxConnection extends EventTarget {
   }
 
   #removeSubscription(id: string) {
+    this.#setSubscriptionConnState(id, FluxSubscriptionState.CLOSED);
+    this.#statusOfSubscriptions.delete(id);
+    this.#onStateChangeCallbacks.delete(id);
     this.#onNextCallbacks.delete(id);
     this.#onCompleteCallbacks.delete(id);
     this.#onErrorCallbacks.delete(id);

diff --git a/packages/ts/frontend/test/FluxConnection.test.ts b/packages/ts/frontend/test/FluxConnection.test.ts
@@ -1,6 +1,7 @@
 import { expect } from '@esm-bundle/chai';
 import type { ReactiveController, ReactiveControllerHost } from '@lit/reactive-element';
 import sinon from 'sinon';
+import { FluxSubscriptionState } from '../FluxConnection';
 import { ActionOnLostSubscription, FluxConnection, State } from '../src/FluxConnection.js';
 import type {
   AbstractMessage,
@@ -80,7 +81,7 @@ describe('@vaadin/hilla-frontend', () => {
 
     it('should immediately return a Subscription when subscribing', () => {
       const sub = fluxConnection.subscribe('MyEndpoint', 'myMethod');
-      for (const name of ['onNext', 'onComplete', 'onError']) {
+      for (const name of ['onNext', 'onComplete', 'onError', 'onConnectionStateChange', 'onSubscriptionLost']) {
         expect(sub).to.have.property(name).which.is.a('function');
       }
     });
@@ -373,7 +374,7 @@ describe('@vaadin/hilla-frontend', () => {
       sub.onSubscriptionLost(resubscribe);
       getSubscriptionEventSpies()?.onReconnect?.();
       getSubscriptionEventSpies()?.onReopen?.();
-      expect(resubscribe).to.have.been.calledOnce;
+      expect(resubscribe).to.have.been.calledTwice;
       expect(getSubscriptionEventSpies()?.push).to.have.been.calledWith(
         JSON.stringify({
           '@type': 'subscribe',
@@ -394,7 +395,7 @@ describe('@vaadin/hilla-frontend', () => {
       getSubscriptionEventSpies()?.onReopen?.();
       getSubscriptionEventSpies()?.onReconnect?.();
       getSubscriptionEventSpies()?.onReopen?.();
-      expect(resubscribe).to.have.been.calledOnce;
+      expect(resubscribe).to.have.been.calledTwice;
     });
 
     it('should remove subscription information when callback has no return value', () => {
@@ -405,7 +406,59 @@ describe('@vaadin/hilla-frontend', () => {
       getSubscriptionEventSpies()?.onReopen?.();
       getSubscriptionEventSpies()?.onReconnect?.();
       getSubscriptionEventSpies()?.onReopen?.();
-      expect(resubscribe).to.have.been.calledOnce;
+      expect(resubscribe).to.have.been.calledTwice;
+    });
+
+    describe('flux subscription state', () => {
+      it('calls change subscription connection state callback', () => {
+        const sub = fluxConnection.subscribe('MyEndpoint', 'myMethod');
+        let subState;
+        sub.onConnectionStateChange((event) => {
+          subState = event.detail.state;
+        });
+        expect(subState).to.be.eq(FluxSubscriptionState.CONNECTING);
+        emitMessage({ '@type': 'update', id: '0' });
+        expect(subState).to.be.eq(FluxSubscriptionState.CONNECTED);
+        getSubscriptionEventSpies()?.onClose?.();
+        expect(subState).to.be.eq(FluxSubscriptionState.CONNECTED);
+        getSubscriptionEventSpies()?.onReconnect?.();
+        expect(subState).to.be.eq(FluxSubscriptionState.CLOSED);
+        getSubscriptionEventSpies()?.onReopen?.();
+        expect(subState).to.be.eq(FluxSubscriptionState.CLOSED);
+      });
+
+      it('calls change subscription connection state callback with re-subscription', () => {
+        const sub = fluxConnection.subscribe('MyEndpoint', 'myMethod');
+        let subState;
+        const resubscribe = sinon.stub();
+        resubscribe.returns(ActionOnLostSubscription.RESUBSCRIBE);
+        sub.onSubscriptionLost(resubscribe);
+        sub.onConnectionStateChange((event) => {
+          subState = event.detail.state;
+        });
+        expect(subState).to.be.eq(FluxSubscriptionState.CONNECTING);
+        emitMessage({ '@type': 'update', id: '0' });
+        expect(subState).to.be.eq(FluxSubscriptionState.CONNECTED);
+        getSubscriptionEventSpies()?.onClose?.();
+        expect(subState).to.be.eq(FluxSubscriptionState.CONNECTED);
+        getSubscriptionEventSpies()?.onReconnect?.();
+        expect(subState).to.be.eq(FluxSubscriptionState.CONNECTING);
+        getSubscriptionEventSpies()?.onReopen?.();
+        expect(subState).to.be.eq(FluxSubscriptionState.CONNECTING);
+        emitMessage({ '@type': 'update', id: '0' });
+        expect(subState).to.be.eq(FluxSubscriptionState.CONNECTED);
+      });
+
+      it('should close subscription connection state if failed to reconnect', () => {
+        const sub = fluxConnection.subscribe('MyEndpoint', 'myMethod');
+        let subState;
+        sub.onConnectionStateChange((event) => {
+          subState = event.detail.state;
+        });
+        expect(subState).to.be.eq(FluxSubscriptionState.CONNECTING);
+        getSubscriptionEventSpies()?.onFailureToReconnect?.();
+        expect(subState).to.be.eq(FluxSubscriptionState.CLOSED);
+      });
     });
   });
 });
diff --git a/packages/ts/react-signals/test/FullStackSignal.spec.tsx b/packages/ts/react-signals/test/FullStackSignal.spec.tsx
@@ -112,6 +112,9 @@ describe('@vaadin/hilla-react-signals', () => {
         onSubscriptionLost() {
           return this;
         },
+        onConnectionStateChange() {
+          return this;
+        },
       });
       // Mock the subscribe method
       client.subscribe.returns(subscription);

diff --git a/packages/ts/react-signals/test/utils.ts b/packages/ts/react-signals/test/utils.ts
@@ -30,6 +30,9 @@ export function createSubscriptionStub<T>(): sinon.SinonSpiedInstance<Subscripti
     onSubscriptionLost() {
       return this;
     },
+    onConnectionStateChange() {
+      return this;
+    },
   });
 }