DO NOT MERGE - draft PR for custom session store strategy

packages/core/src/domain/configuration/configuration.spec.ts

@@ -7,6 +7,8 @@ import {
   isExperimentalFeatureEnabled,
   resetExperimentalFeatures,
 } from '../../tools/experimentalFeatures'
+import { noop } from '../../tools/utils/functionUtils'
+import type { SessionStoreStrategy, SessionStoreStrategyMethod } from '../session/storeStrategies/sessionStoreStrategy'
 import { TrackingConsent } from '../trackingConsent'
 import type { InitConfiguration } from './configuration'
 import { serializeConfiguration, validateAndBuildConfiguration } from './configuration'
@@ -99,12 +101,30 @@ describe('validateAndBuildConfiguration', () => {
   })
 
   describe('sessionStoreStrategyType', () => {
+    const customSessionStoreStrategy: SessionStoreStrategy = {
+      expireSession: noop,
+      isLockEnabled: false,
+      retrieveSession: () => ({}),
+      persistSession: noop,
+    }
+
     it('allowFallbackToLocalStorage should not be enabled by default', () => {
       spyOnProperty(document, 'cookie', 'get').and.returnValue('')
       const configuration = validateAndBuildConfiguration({ clientToken })
       expect(configuration?.sessionStoreStrategyType).toBeUndefined()
     })
 
+    it('should use custom strategy if provided, even when cookies are available', () => {
+      const configuration = validateAndBuildConfiguration({ clientToken, customSessionStoreStrategy })
+      expect(configuration?.sessionStoreStrategyType?.type).toBe('Custom')
+      if (configuration?.sessionStoreStrategyType?.type === 'Custom') {
+        expect(configuration?.sessionStoreStrategyType?.sessionStoreStrategy.expireSession).toBeInstanceOf(Function)
+        expect(configuration?.sessionStoreStrategyType?.sessionStoreStrategy.retrieveSession).toBeInstanceOf(Function)
+        expect(configuration?.sessionStoreStrategyType?.sessionStoreStrategy.persistSession).toBeInstanceOf(Function)
+        expect(configuration?.sessionStoreStrategyType?.sessionStoreStrategy.isLockEnabled).toBe(false)
+      }
+    })
+
     it('should contain cookie in the configuration by default', () => {
       const configuration = validateAndBuildConfiguration({ clientToken, allowFallbackToLocalStorage: false })
       expect(configuration?.sessionStoreStrategyType).toEqual({
@@ -133,6 +153,23 @@ describe('validateAndBuildConfiguration', () => {
       const configuration = validateAndBuildConfiguration({ clientToken, allowFallbackToLocalStorage: true })
       expect(configuration?.sessionStoreStrategyType).toBeUndefined()
     })
+
+    it('should contain custom strategy when present if both cookies and local storage are unavailable', () => {
+      spyOnProperty(document, 'cookie', 'get').and.returnValue('')
+      spyOn(Storage.prototype, 'getItem').and.throwError('unavailable')
+      const configuration = validateAndBuildConfiguration({
+        clientToken,
+        allowFallbackToLocalStorage: true,
+        customSessionStoreStrategy,
+      })
+      expect(configuration?.sessionStoreStrategyType?.type).toBe('Custom')
+      if (configuration?.sessionStoreStrategyType?.type === 'Custom') {
+        expect(configuration?.sessionStoreStrategyType?.sessionStoreStrategy.expireSession).toBeInstanceOf(Function)
+        expect(configuration?.sessionStoreStrategyType?.sessionStoreStrategy.retrieveSession).toBeInstanceOf(Function)
+        expect(configuration?.sessionStoreStrategyType?.sessionStoreStrategy.persistSession).toBeInstanceOf(Function)
+        expect(configuration?.sessionStoreStrategyType?.sessionStoreStrategy.isLockEnabled).toBe(false)
+      }
+    })
   })
 
   describe('beforeSend', () => {
@@ -241,6 +278,69 @@ describe('validateAndBuildConfiguration', () => {
     })
   })
 
+  describe('customSessionStoreStrategy validation', () => {
+    const baseCustomSessionStoreStrategy: SessionStoreStrategy = {
+      expireSession: noop,
+      persistSession: noop,
+      retrieveSession: () => ({}),
+      isLockEnabled: false,
+    }
+
+    ;(['expireSession', 'persistSession', 'retrieveSession'] satisfies SessionStoreStrategyMethod[]).forEach(
+      (method) => {
+        it(`should validate the ${method} method`, () => {
+          validateAndBuildConfiguration({
+            clientToken,
+            customSessionStoreStrategy: {
+              ...baseCustomSessionStoreStrategy,
+              [method]: true,
+            },
+          })
+          expect(displaySpy).toHaveBeenCalledOnceWith(`customSessionStoreStrategy.${method} should be a function`)
+        })
+      }
+    )
+
+    it('should wrap all methods on the strategy for error handling', () => {
+      const error = new Error('Method failed')
+      const customSessionStoreStrategy: SessionStoreStrategy = {
+        isLockEnabled: true,
+        expireSession: () => {
+          throw error
+        },
+        retrieveSession: () => {
+          throw error
+          // eslint-disable-next-line no-unreachable
+          return {}
+        },
+        persistSession: () => {
+          throw error
+        },
+      }
+
+      const expireSpy = spyOn(customSessionStoreStrategy, 'expireSession').and.callThrough()
+      const retrieveSpy = spyOn(customSessionStoreStrategy, 'retrieveSession').and.callThrough()
+      const persistSpy = spyOn(customSessionStoreStrategy, 'persistSession').and.callThrough()
+
+      const configuration = validateAndBuildConfiguration({ clientToken, customSessionStoreStrategy })
+      expect(configuration?.customSessionStoreStrategy).toBeDefined()
+      expect(configuration?.customSessionStoreStrategy).not.toBe(customSessionStoreStrategy)
+      expect(configuration?.customSessionStoreStrategy).not.toEqual(customSessionStoreStrategy)
+
+      expect(configuration?.customSessionStoreStrategy?.expireSession()).toBeUndefined()
+      expect(displaySpy).toHaveBeenCalledWith('customSessionStoryStrategy.expireSession threw an error', error)
+      expect(expireSpy).toHaveBeenCalled()
+
+      expect(configuration?.customSessionStoreStrategy?.persistSession({})).toBeUndefined()
+      expect(displaySpy).toHaveBeenCalledWith('customSessionStoryStrategy.persistSession threw an error', error)
+      expect(persistSpy).toHaveBeenCalledWith({})
+
+      expect(configuration?.customSessionStoreStrategy?.retrieveSession()).toEqual({})
+      expect(displaySpy).toHaveBeenCalledWith('customSessionStoryStrategy.retrieveSession threw an error', error)
+      expect(retrieveSpy).toHaveBeenCalled()
+    })
+  })
+
   describe('serializeConfiguration', () => {
     it('should serialize the configuration', () => {
       // By specifying the type here, we can ensure that serializeConfiguration is returning an

packages/core/src/domain/configuration/configuration.ts

@@ -8,7 +8,11 @@ import { ONE_KIBI_BYTE } from '../../tools/utils/byteUtils'
 import { objectHasValue } from '../../tools/utils/objectUtils'
 import { assign } from '../../tools/utils/polyfills'
 import { selectSessionStoreStrategyType } from '../session/sessionStore'
-import type { SessionStoreStrategyType } from '../session/storeStrategies/sessionStoreStrategy'
+import type {
+  SessionStoreStrategy,
+  SessionStoreStrategyMethod,
+  SessionStoreStrategyType,
+} from '../session/storeStrategies/sessionStoreStrategy'
 import { TrackingConsent } from '../trackingConsent'
 import type { TransportConfiguration } from './transportConfiguration'
 import { computeTransportConfiguration } from './transportConfiguration'
@@ -59,6 +63,10 @@ export interface InitConfiguration {
    * @default false
    */
   allowUntrustedEvents?: boolean | undefined
+  /**
+   * Pass a custom {@link SessionStoreStrategy} for managing RUM session lifecycle.
+   */
+  customSessionStoreStrategy?: SessionStoreStrategy | undefined
   /**
    * Store global context and user context in localStorage to preserve them along the user navigation.
    * See [Contexts life cycle](https://docs.datadoghq.com/real_user_monitoring/browser/advanced_configuration/?tab=npm#contexts-life-cycle) for further information.
@@ -169,6 +177,7 @@ interface ReplicaUserConfiguration {
 export interface Configuration extends TransportConfiguration {
   // Built from init configuration
   beforeSend: GenericBeforeSendCallback | undefined
+  customSessionStoreStrategy: SessionStoreStrategy | undefined
   sessionStoreStrategyType: SessionStoreStrategyType | undefined
   sessionSampleRate: number
   telemetrySampleRate: number
@@ -215,6 +224,15 @@ export function isSampleRate(sampleRate: unknown, name: string) {
   return true
 }
 
+export function isFunction(value: unknown, path: string) {
+  if (!value || typeof value !== 'function') {
+    display.error(`${path} should be a function`)
+    return false
+  }
+
+  return true
+}
+
 export function validateAndBuildConfiguration(initConfiguration: InitConfiguration): Configuration | undefined {
   if (!initConfiguration || !initConfiguration.clientToken) {
     display.error('Client Token is not configured, we will not send any data.')
@@ -242,10 +260,46 @@ export function validateAndBuildConfiguration(initConfiguration: InitConfigurati
     return
   }
 
+  let wrappedCustomSessionStoreStrategy: SessionStoreStrategy | undefined
+
+  if (initConfiguration.customSessionStoreStrategy !== undefined) {
+    const strategyMethods = [
+      'expireSession',
+      'persistSession',
+      'retrieveSession',
+    ] satisfies SessionStoreStrategyMethod[]
+    if (
+      !strategyMethods.every((methodName) =>
+        isFunction(
+          initConfiguration.customSessionStoreStrategy?.[methodName],
+          `customSessionStoreStrategy.${methodName}`
+        )
+      )
+    ) {
+      return
+    }
+
+    const { isLockEnabled, expireSession, persistSession, retrieveSession } =
+      initConfiguration.customSessionStoreStrategy
+
+    const wrappedRetrieveSession = catchUserErrors(
+      retrieveSession,
+      'customSessionStoryStrategy.retrieveSession threw an error'
+    )
+
+    wrappedCustomSessionStoreStrategy = {
+      expireSession: catchUserErrors(expireSession, 'customSessionStoryStrategy.expireSession threw an error'),
+      isLockEnabled,
+      persistSession: catchUserErrors(persistSession, 'customSessionStoryStrategy.persistSession threw an error'),
+      retrieveSession: () => wrappedRetrieveSession() ?? {},
+    }
+  }
+
   return assign(
     {
       beforeSend:
         initConfiguration.beforeSend && catchUserErrors(initConfiguration.beforeSend, 'beforeSend threw an error:'),
+      customSessionStoreStrategy: wrappedCustomSessionStoreStrategy,
       sessionStoreStrategyType: selectSessionStoreStrategyType(initConfiguration),
       sessionSampleRate: initConfiguration.sessionSampleRate ?? 100,
       telemetrySampleRate: initConfiguration.telemetrySampleRate ?? 20,

packages/core/src/domain/session/sessionStore.ts

@@ -5,6 +5,7 @@ import { throttle } from '../../tools/utils/functionUtils'
 import { generateUUID } from '../../tools/utils/stringUtils'
 import type { InitConfiguration } from '../configuration'
 import { assign } from '../../tools/utils/polyfills'
+import { display } from '../../tools/display'
 import { selectCookieStrategy, initCookieStrategy } from './storeStrategies/sessionInCookie'
 import type { SessionStoreStrategyType } from './storeStrategies/sessionStoreStrategy'
 import {
@@ -38,12 +39,22 @@ export interface SessionStore {
 export const STORAGE_POLL_DELAY = ONE_SECOND
 
 /**
- * Checks if cookies are available as the preferred storage
- * Else, checks if LocalStorage is allowed and available
+ * Returns:
+ * 1. Custom session storage strategy, if available
+ * 2. Cookies strategy, if available and preferred
+ * 3. LocalStorage strategy if allowed and available
  */
 export function selectSessionStoreStrategyType(
   initConfiguration: InitConfiguration
 ): SessionStoreStrategyType | undefined {
+  if (initConfiguration.customSessionStoreStrategy) {
+    display.warn('Using custom session store strategy.')
+    return {
+      type: 'Custom',
+      sessionStoreStrategy: initConfiguration.customSessionStoreStrategy,
+    }
+  }
+
   let sessionStoreStrategyType = selectCookieStrategy(initConfiguration)
   if (!sessionStoreStrategyType && initConfiguration.allowFallbackToLocalStorage) {
     sessionStoreStrategyType = selectLocalStorageStrategy()
@@ -67,9 +78,11 @@ export function startSessionStore<TrackingType extends string>(
   const sessionStateUpdateObservable = new Observable<{ previousState: SessionState; newState: SessionState }>()
 
   const sessionStoreStrategy =
-    sessionStoreStrategyType.type === 'Cookie'
-      ? initCookieStrategy(sessionStoreStrategyType.cookieOptions)
-      : initLocalStorageStrategy()
+    sessionStoreStrategyType.type === 'Custom'
+      ? sessionStoreStrategyType.sessionStoreStrategy
+      : sessionStoreStrategyType.type === 'Cookie'
+        ? initCookieStrategy(sessionStoreStrategyType.cookieOptions)
+        : initLocalStorageStrategy()
   const { expireSession } = sessionStoreStrategy
 
   const watchSessionTimeoutId = setInterval(watchSession, STORAGE_POLL_DELAY)

packages/core/src/domain/session/storeStrategies/sessionStoreStrategy.ts

@@ -3,11 +3,18 @@ import type { SessionState } from '../sessionState'
 
 export const SESSION_STORE_KEY = '_dd_s'
 
-export type SessionStoreStrategyType = { type: 'Cookie'; cookieOptions: CookieOptions } | { type: 'LocalStorage' }
+export type SessionStoreStrategyType =
+  | { type: 'Cookie'; cookieOptions: CookieOptions }
+  | { type: 'LocalStorage' }
+  | { type: 'Custom'; sessionStoreStrategy: SessionStoreStrategy }
 
 export interface SessionStoreStrategy {
   isLockEnabled: boolean
   persistSession: (session: SessionState) => void
   retrieveSession: () => SessionState
   expireSession: () => void
 }
+
+export type SessionStoreStrategyMethod = keyof {
+  [K in keyof SessionStoreStrategy as SessionStoreStrategy[K] extends (...args: any[]) => unknown ? K : never]: K
+}

packages/core/src/index.ts

@@ -123,7 +123,7 @@ export { CustomerDataType } from './domain/context/contextConstants'
 export { createValueHistory, ValueHistory, ValueHistoryEntry, CLEAR_OLD_VALUES_INTERVAL } from './tools/valueHistory'
 export { readBytesFromStream } from './tools/readBytesFromStream'
 export { STORAGE_POLL_DELAY } from './domain/session/sessionStore'
-export { SESSION_STORE_KEY } from './domain/session/storeStrategies/sessionStoreStrategy'
+export { SESSION_STORE_KEY, SessionStoreStrategy } from './domain/session/storeStrategies/sessionStoreStrategy'
 export {
   willSyntheticsInjectRum,
   getSyntheticsTestId,

packages/core/test/coreConfiguration.ts

@@ -1,5 +1,7 @@
 import type { InitConfiguration } from '../src/domain/configuration'
+import type { SessionState } from '../src/domain/session/sessionState'
 import type { RawTelemetryConfiguration } from '../src/domain/telemetry'
+import { noop } from '../src/tools/utils/functionUtils'
 import type { CamelToSnakeCase, RemoveIndex } from './typeUtils'
 
 // Defines a few constants and types related to the core package configuration, so it can be used in
@@ -12,6 +14,12 @@ import type { CamelToSnakeCase, RemoveIndex } from './typeUtils'
 export const EXHAUSTIVE_INIT_CONFIGURATION: Required<InitConfiguration> = {
   clientToken: 'yes',
   beforeSend: () => true,
+  customSessionStoreStrategy: {
+    expireSession: noop,
+    isLockEnabled: false,
+    persistSession: noop,
+    retrieveSession: () => ({}) satisfies SessionState,
+  },
   sessionSampleRate: 50,
   telemetrySampleRate: 60,
   silentMultipleInit: true,
@@ -74,6 +82,8 @@ export type MapInitConfigurationKey<Key extends string> =
           | 'internalAnalyticsSubdomain'
           | 'replica'
           | 'enableExperimentalFeatures'
+          // TODO: Convert to a flag like `beforeSend` and include in telemetry
+          | 'customSessionStoreStrategy'
       ? never
       : // Other keys are simply snake cased
         CamelToSnakeCase<Key>

packages/rum-core/src/domain/configuration/configuration.ts

@@ -195,6 +195,7 @@ export function validateAndBuildRumConfiguration(
       actionNameAttribute: initConfiguration.actionNameAttribute,
       sessionReplaySampleRate: initConfiguration.sessionReplaySampleRate ?? 0,
       startSessionReplayRecordingManually: !!initConfiguration.startSessionReplayRecordingManually,
+      customSessionStoreStrategy: initConfiguration.customSessionStoreStrategy,
       traceSampleRate: initConfiguration.traceSampleRate,
       allowedTracingUrls,
       excludedActivityUrls: initConfiguration.excludedActivityUrls ?? [],

packages/rum/src/entries/main.ts

@@ -29,7 +29,7 @@ export {
   RumOtherResourceEventDomainContext,
   RumLongTaskEventDomainContext,
 } from '@datadog/browser-rum-core'
-export { DefaultPrivacyLevel } from '@datadog/browser-core'
+export { DefaultPrivacyLevel, SessionStoreStrategy } from '@datadog/browser-core'
 
 const recorderApi = makeRecorderApi(startRecording)
 export const datadogRum = makeRumPublicApi(startRum, recorderApi, { startDeflateWorker, createDeflateEncoder })