feat: Move deregisterDrain method to feature-base (#1220)

Co-authored-by: Chunwai Li <[email protected]>

Author: ellisong

src/features/generic_events/aggregate/index.js

@@ -11,7 +11,6 @@ import { AggregateBase } from '../../utils/aggregate-base'
 import { warn } from '../../../common/util/console'
 import { now } from '../../../common/timing/now'
 import { registerHandler } from '../../../common/event-emitter/register-handler'
-import { deregisterDrain } from '../../../common/drain/drain'
 import { SUPPORTABILITY_METRIC_CHANNEL } from '../../metrics/constants'
 import { EventBuffer } from '../../utils/event-buffer'
 import { applyFnToProps } from '../../../common/util/traverse'
@@ -33,7 +32,7 @@ export class Aggregate extends AggregateBase {
     this.waitForFlags(['ins']).then(([ins]) => {
       if (!ins) {
         this.blocked = true
-        deregisterDrain(this.agentIdentifier, this.featureName)
+        this.deregisterDrain()
         return
       }
 

src/features/generic_events/instrument/index.js

@@ -3,7 +3,6 @@
  */
 
 import { isBrowserScope } from '../../../common/constants/runtime'
-import { deregisterDrain } from '../../../common/drain/drain'
 import { handle } from '../../../common/event-emitter/handle'
 import { windowAddEventListener } from '../../../common/event-listener/event-listener-opts'
 import { InstrumentBase } from '../../utils/instrument-base'
@@ -31,7 +30,7 @@ export class Instrument extends InstrumentBase {
 
     /** If any of the sources are active, import the aggregator. otherwise deregister */
     if (genericEventSourceConfigs.some(x => x)) this.importAggregator(agentRef)
-    else deregisterDrain(agentRef.agentIdentifier, this.featureName)
+    else this.deregisterDrain()
   }
 }
 

src/features/jserrors/aggregate/index.js

@@ -18,7 +18,6 @@ import { FEATURE_NAME } from '../constants'
 import { FEATURE_NAMES } from '../../../loaders/features/features'
 import { AggregateBase } from '../../utils/aggregate-base'
 import { getNREUMInitializedAgent } from '../../../common/window/nreum'
-import { deregisterDrain } from '../../../common/drain/drain'
 import { now } from '../../../common/timing/now'
 import { applyFnToProps } from '../../../common/util/traverse'
 import { evaluateInternalError } from './internal-errors'
@@ -59,7 +58,7 @@ export class Aggregate extends AggregateBase {
         this.drain()
       } else {
         this.blocked = true // if rum response determines that customer lacks entitlements for spa endpoint, this feature shouldn't harvest
-        deregisterDrain(this.agentIdentifier, this.featureName)
+        this.deregisterDrain()
       }
     })
   }

src/features/metrics/aggregate/index.js

@@ -7,7 +7,6 @@ import { onDOMContentLoaded } from '../../../common/window/load'
 import { windowAddEventListener } from '../../../common/event-listener/event-listener-opts'
 import { isBrowserScope, isWorkerScope } from '../../../common/constants/runtime'
 import { AggregateBase } from '../../utils/aggregate-base'
-import { deregisterDrain } from '../../../common/drain/drain'
 import { isIFrameWindow } from '../../../common/dom/iframe'
 // import { WEBSOCKET_TAG } from '../../../common/wrap/wrap-websocket'
 // import { handleWebsocketEvents } from './websocket-detection'
@@ -26,7 +25,7 @@ export class Aggregate extends AggregateBase {
         this.drain()
       } else {
         this.blocked = true // if rum response determines that customer lacks entitlements for spa endpoint, this feature shouldn't harvest
-        deregisterDrain(this.agentIdentifier, this.featureName)
+        this.deregisterDrain()
       }
     })
 

src/features/session_replay/aggregate/index.js

@@ -21,7 +21,6 @@ import { RRWEB_VERSION } from '../../../common/constants/env'
 import { MODE, SESSION_EVENTS, SESSION_EVENT_TYPES } from '../../../common/session/constants'
 import { stringify } from '../../../common/util/stringify'
 import { stylesheetEvaluator } from '../shared/stylesheet-evaluator'
-import { deregisterDrain } from '../../../common/drain/drain'
 import { now } from '../../../common/timing/now'
 import { buildNRMetaNode } from '../shared/utils'
 import { MAX_PAYLOAD_SIZE } from '../../../common/constants/agent-constants'
@@ -106,7 +105,7 @@ export class Aggregate extends AggregateBase {
     this.waitForFlags(['srs', 'sr']).then(([srMode, entitled]) => {
       this.entitled = !!entitled
       if (!this.entitled) {
-        deregisterDrain(this.agentIdentifier, this.featureName)
+        this.deregisterDrain()
         if (this.recorder?.recording) {
           this.abort(ABORT_REASONS.ENTITLEMENTS)
           handle(SUPPORTABILITY_METRIC_CHANNEL, ['SessionReplay/EnabledNotEntitled/Detected'], undefined, FEATURE_NAMES.metrics, this.ee)

src/features/session_trace/aggregate/index.js

@@ -4,7 +4,6 @@ import { FEATURE_NAME } from '../constants'
 import { AggregateBase } from '../../utils/aggregate-base'
 import { TraceStorage } from './trace/storage'
 import { obj as encodeObj } from '../../../common/url/encode'
-import { deregisterDrain } from '../../../common/drain/drain'
 import { globalScope } from '../../../common/constants/runtime'
 import { MODE, SESSION_EVENTS } from '../../../common/session/constants'
 import { applyFnToProps } from '../../../common/util/traverse'
@@ -38,7 +37,7 @@ export class Aggregate extends AggregateBase {
   /** Sets up event listeners, and initializes this module to run in the correct "mode".  Can be triggered from a few places, but makes an effort to only set up listeners once */
   initialize (stMode, stEntitled, ignoreSession) {
     this.entitled ??= stEntitled
-    if (this.blocked || !this.entitled) return deregisterDrain(this.agentIdentifier, this.featureName)
+    if (this.blocked || !this.entitled) return this.deregisterDrain()
 
     if (!this.initialized) {
       this.initialized = true
@@ -74,7 +73,7 @@ export class Aggregate extends AggregateBase {
 
     /** If the mode is off, we do not want to hold up draining for other features, so we deregister the feature for now.
      * If it drains later (due to a mode change), data and handlers will instantly drain instead of waiting for the registry. */
-    if (this.mode === MODE.OFF) return deregisterDrain(this.agentIdentifier, this.featureName)
+    if (this.mode === MODE.OFF) return this.deregisterDrain()
 
     this.timeKeeper ??= this.agentRef.runtime.timeKeeper
 

src/features/session_trace/instrument/index.js

@@ -8,7 +8,6 @@ import { wrapEvents } from '../../../common/wrap/wrap-events'
 import { InstrumentBase } from '../../utils/instrument-base'
 import * as CONSTANTS from '../constants'
 import { FEATURE_NAMES } from '../../../loaders/features/features'
-import { deregisterDrain } from '../../../common/drain/drain'
 import { canEnableSessionTracking } from '../../utils/feature-gates'
 import { now } from '../../../common/timing/now'
 
@@ -22,7 +21,7 @@ export class Instrument extends InstrumentBase {
     super(agentRef, FEATURE_NAME, auto)
     const canTrackSession = canEnableSessionTracking(this.agentIdentifier)
     if (!canTrackSession) {
-      deregisterDrain(this.agentIdentifier, this.featureName)
+      this.deregisterDrain()
       return
     }
 

src/features/soft_navigations/aggregate/index.js

@@ -1,4 +1,3 @@
-import { deregisterDrain } from '../../../common/drain/drain'
 import { handle } from '../../../common/event-emitter/handle'
 import { registerHandler } from '../../../common/event-emitter/register-handler'
 import { HarvestScheduler } from '../../../common/harvest/harvest-scheduler'
@@ -49,7 +48,7 @@ export class Aggregate extends AggregateBase {
         scheduler.startTimer(harvestTimeSeconds, 0)
       } else {
         this.blocked = true // if rum response determines that customer lacks entitlements for spa endpoint, this feature shouldn't harvest
-        deregisterDrain(this.agentIdentifier, this.featureName)
+        this.deregisterDrain()
       }
     })
 

src/features/spa/aggregate/index.js

@@ -22,7 +22,6 @@ import { bundleId } from '../../../common/ids/bundle-id'
 import { loadedAsDeferredBrowserScript } from '../../../common/constants/runtime'
 import { handle } from '../../../common/event-emitter/handle'
 import { SUPPORTABILITY_METRIC_CHANNEL } from '../../metrics/constants'
-import { deregisterDrain } from '../../../common/drain/drain'
 import { warn } from '../../../common/util/console'
 import { EventBuffer } from '../../utils/event-buffer'
 
@@ -112,7 +111,7 @@ export class Aggregate extends AggregateBase {
         this.drain()
       } else {
         this.blocked = true
-        deregisterDrain(agentRef.agentIdentifier, this.featureName)
+        this.deregisterDrain()
       }
     })
 

src/features/utils/aggregate-base.js

@@ -2,7 +2,7 @@ import { FeatureBase } from './feature-base'
 import { isValid } from '../../common/config/info'
 import { configure } from '../../loaders/configure/configure'
 import { gosCDN } from '../../common/window/nreum'
-import { deregisterDrain, drain } from '../../common/drain/drain'
+import { drain } from '../../common/drain/drain'
 import { activatedFeatures } from '../../common/util/feature-flags'
 import { Obfuscator } from '../../common/util/obfuscate'
 
@@ -38,7 +38,7 @@ export class AggregateBase extends FeatureBase {
     return flagsPromise.catch(err => {
       this.ee.emit('internal-error', [err])
       this.blocked = true
-      deregisterDrain(this.agentIdentifier, this.featureName)
+      this.deregisterDrain()
     })
   }
 

src/features/utils/feature-base.js

@@ -1,4 +1,5 @@
 import { ee } from '../../common/event-emitter/contextual-ee'
+import { deregisterDrain } from '../../common/drain/drain'
 
 export class FeatureBase {
   constructor (agentIdentifier, featureName) {
@@ -15,4 +16,8 @@ export class FeatureBase {
      */
     this.blocked = false
   }
+
+  deregisterDrain () {
+    deregisterDrain(this.agentIdentifier, this.featureName)
+  }
 }