fix: Ensure streaming connection is closed on SDK close. (#774)

Author: kinyoklion

packages/sdk/browser/__tests__/BrowserDataManager.test.ts

@@ -63,13 +63,16 @@ describe('given a BrowserDataManager with mocked dependencies', () => {
   let diagnosticsManager: jest.Mocked<internal.DiagnosticsManager>;
   let dataManager: BrowserDataManager;
   let logger: LDLogger;
+  let eventSourceCloseMethod: jest.Mock;
+
   beforeEach(() => {
     logger = {
       error: jest.fn(),
       warn: jest.fn(),
       info: jest.fn(),
       debug: jest.fn(),
     };
+    eventSourceCloseMethod = jest.fn();
     config = {
       logger,
       maxCachedContexts: 5,
@@ -106,7 +109,7 @@ describe('given a BrowserDataManager with mocked dependencies', () => {
           options,
           onclose: jest.fn(),
           addEventListener: jest.fn(),
-          close: jest.fn(),
+          close: eventSourceCloseMethod,
         })),
         fetch: mockedFetch,
         getEventSourceCapabilities: jest.fn(),
@@ -495,4 +498,31 @@ describe('given a BrowserDataManager with mocked dependencies', () => {
 
     expect(platform.requests.createEventSource).toHaveBeenCalled();
   });
+
+  it('closes the event source when the data manager is closed', async () => {
+    const context = Context.fromLDContext({ kind: 'user', key: 'test-user' });
+    const identifyOptions: BrowserIdentifyOptions = {};
+    const identifyResolve = jest.fn();
+    const identifyReject = jest.fn();
+
+    dataManager.setForcedStreaming(undefined);
+    dataManager.setAutomaticStreamingState(true);
+    expect(platform.requests.createEventSource).not.toHaveBeenCalled();
+
+    flagManager.loadCached.mockResolvedValue(false);
+
+    await dataManager.identify(identifyResolve, identifyReject, context, identifyOptions);
+
+    expect(platform.requests.createEventSource).toHaveBeenCalled();
+
+    dataManager.close();
+    expect(eventSourceCloseMethod).toHaveBeenCalled();
+    // Verify a subsequent identify doesn't create a new event source
+    await dataManager.identify(identifyResolve, identifyReject, context, {});
+    expect(platform.requests.createEventSource).toHaveBeenCalledTimes(1);
+
+    expect(logger.debug).toHaveBeenCalledWith(
+      '[BrowserDataManager] Identify called after data manager was closed.',
+    );
+  });
 });

packages/sdk/browser/src/BrowserDataManager.ts

@@ -74,6 +74,11 @@ export default class BrowserDataManager extends BaseDataManager {
     context: Context,
     identifyOptions?: LDIdentifyOptions,
   ): Promise<void> {
+    if (this.closed) {
+      this._debugLog('Identify called after data manager was closed.');
+      return;
+    }
+
     this.context = context;
     const browserIdentifyOptions = identifyOptions as BrowserIdentifyOptions | undefined;
     if (browserIdentifyOptions?.hash) {

packages/sdk/react-native/__tests__/MobileDataManager.test.ts

@@ -54,13 +54,17 @@ describe('given a MobileDataManager with mocked dependencies', () => {
   let diagnosticsManager: jest.Mocked<internal.DiagnosticsManager>;
   let mobileDataManager: MobileDataManager;
   let logger: LDLogger;
+  let eventSourceCloseMethod: jest.Mock;
+
   beforeEach(() => {
     logger = {
       error: jest.fn(),
       warn: jest.fn(),
       info: jest.fn(),
       debug: jest.fn(),
     };
+    eventSourceCloseMethod = jest.fn();
+
     config = {
       logger,
       maxCachedContexts: 5,
@@ -94,7 +98,7 @@ describe('given a MobileDataManager with mocked dependencies', () => {
           options,
           onclose: jest.fn(),
           addEventListener: jest.fn(),
-          close: jest.fn(),
+          close: eventSourceCloseMethod,
         })),
         fetch: mockedFetch,
         getEventSourceCapabilities: jest.fn(),
@@ -223,6 +227,23 @@ describe('given a MobileDataManager with mocked dependencies', () => {
     expect(platform.requests.fetch).not.toHaveBeenCalled();
   });
 
+  it('makes no connection when closed', async () => {
+    mobileDataManager.close();
+
+    const context = Context.fromLDContext({ kind: 'user', key: 'test-user' });
+    const identifyOptions: LDIdentifyOptions = { waitForNetworkResults: false };
+    const identifyResolve = jest.fn();
+    const identifyReject = jest.fn();
+
+    await mobileDataManager.identify(identifyResolve, identifyReject, context, identifyOptions);
+
+    expect(platform.requests.createEventSource).not.toHaveBeenCalled();
+    expect(platform.requests.fetch).not.toHaveBeenCalled();
+    expect(logger.debug).toHaveBeenCalledWith(
+      '[MobileDataManager] Identify called after data manager was closed.',
+    );
+  });
+
   it('should load cached flags and resolve the identify', async () => {
     const context = Context.fromLDContext({ kind: 'user', key: 'test-user' });
     const identifyOptions: LDIdentifyOptions = { waitForNetworkResults: false };
@@ -299,4 +320,25 @@ describe('given a MobileDataManager with mocked dependencies', () => {
     expect(identifyResolve).toHaveBeenCalled();
     expect(identifyReject).not.toHaveBeenCalled();
   });
+
+  it('closes the event source when the data manager is closed', async () => {
+    const context = Context.fromLDContext({ kind: 'user', key: 'test-user' });
+    const identifyOptions: LDIdentifyOptions = { waitForNetworkResults: false };
+    const identifyResolve = jest.fn();
+    const identifyReject = jest.fn();
+
+    await mobileDataManager.identify(identifyResolve, identifyReject, context, identifyOptions);
+    expect(platform.requests.createEventSource).toHaveBeenCalled();
+
+    mobileDataManager.close();
+    expect(eventSourceCloseMethod).toHaveBeenCalled();
+
+    // Verify a subsequent identify doesn't create a new event source
+    await mobileDataManager.identify(identifyResolve, identifyReject, context, identifyOptions);
+    expect(platform.requests.createEventSource).toHaveBeenCalledTimes(1);
+
+    expect(logger.debug).toHaveBeenCalledWith(
+      '[MobileDataManager] Identify called after data manager was closed.',
+    );
+  });
 });

packages/sdk/react-native/src/MobileDataManager.ts

@@ -58,6 +58,10 @@ export default class MobileDataManager extends BaseDataManager {
     context: Context,
     identifyOptions?: LDIdentifyOptions,
   ): Promise<void> {
+    if (this.closed) {
+      this._debugLog('Identify called after data manager was closed.');
+      return;
+    }
     this.context = context;
     const offline = this.connectionMode === 'offline';
     // In offline mode we do not support waiting for results.
@@ -140,6 +144,11 @@ export default class MobileDataManager extends BaseDataManager {
   }
 
   async setConnectionMode(mode: ConnectionMode): Promise<void> {
+    if (this.closed) {
+      this._debugLog('setting connection mode after data manager was closed');
+      return;
+    }
+
     if (this.connectionMode === mode) {
       this._debugLog(`setConnectionMode ignored. Mode is already '${mode}'.`);
       return;

packages/shared/sdk-client/__tests__/LDClientImpl.test.ts

@@ -413,4 +413,21 @@ describe('sdk-client object', () => {
       }),
     );
   });
+
+  test('closes event source when client is closed', async () => {
+    const carContext: LDContext = { kind: 'car', key: 'test-car' };
+
+    const mockCreateEventSource = jest.fn((streamUri: string = '', options: any = {}) => {
+      mockEventSource = new MockEventSource(streamUri, options);
+      mockEventSource.simulateEvents('put', [{ data: JSON.stringify(defaultPutResponse) }]);
+      return mockEventSource;
+    });
+    mockPlatform.requests.createEventSource = mockCreateEventSource;
+
+    await ldc.identify(carContext);
+    expect(mockEventSource.closed).toBe(false);
+
+    await ldc.close();
+    expect(mockEventSource.closed).toBe(true);
+  });
 });

packages/shared/sdk-client/src/DataManager.ts

@@ -43,6 +43,11 @@ export interface DataManager {
     context: Context,
     identifyOptions?: LDIdentifyOptions,
   ): Promise<void>;
+
+  /**
+   * Closes the data manager. Any active connections are closed.
+   */
+  close(): void;
 }
 
 /**
@@ -69,6 +74,7 @@ export abstract class BaseDataManager implements DataManager {
   private _connectionParams?: ConnectionParams;
   protected readonly dataSourceStatusManager: DataSourceStatusManager;
   private readonly _dataSourceEventHandler: DataSourceEventHandler;
+  protected closed = false;
 
   constructor(
     protected readonly platform: Platform,
@@ -221,4 +227,9 @@ export abstract class BaseDataManager implements DataManager {
       },
     };
   }
+
+  public close() {
+    this.updateProcessor?.close();
+    this.closed = true;
+  }
 }

packages/shared/sdk-client/src/LDClientImpl.ts

@@ -11,7 +11,6 @@ import {
   LDHeaders,
   LDLogger,
   Platform,
-  subsystem,
   timedPromise,
   TypeValidators,
 } from '@launchdarkly/js-sdk-common';
@@ -48,7 +47,6 @@ export default class LDClientImpl implements LDClient {
   private readonly _diagnosticsManager?: internal.DiagnosticsManager;
   private _eventProcessor?: internal.EventProcessor;
   readonly logger: LDLogger;
-  private _updateProcessor?: subsystem.LDStreamProcessor;
 
   private readonly _highTimeoutThreshold: number = 15;
 
@@ -153,7 +151,7 @@ export default class LDClientImpl implements LDClient {
   async close(): Promise<void> {
     await this.flush();
     this._eventProcessor?.close();
-    this._updateProcessor?.close();
+    this.dataManager.close();
     this.logger.debug('Closed event processor and data source.');
   }