lazy load dsm only when needed (#5305)

rochdev · web-flow · commit bcf7e6129bfc · 2025-02-25T21:52:33.000+01:00
* lazy load dsm only when needed

* code cleanup

* refactor and add comments
diff --git a/packages/datadog-plugin-amqplib/src/consumer.js b/packages/datadog-plugin-amqplib/src/consumer.js
@@ -2,7 +2,7 @@
 
 const { TEXT_MAP } = require('../../../ext/formats')
 const ConsumerPlugin = require('../../dd-trace/src/plugins/consumer')
-const { getAmqpMessageSize } = require('../../dd-trace/src/datastreams/processor')
+const { getAmqpMessageSize } = require('../../dd-trace/src/datastreams')
 const { getResourceName } = require('./util')
 
 class AmqplibConsumerPlugin extends ConsumerPlugin {
diff --git a/packages/datadog-plugin-amqplib/src/producer.js b/packages/datadog-plugin-amqplib/src/producer.js
@@ -3,8 +3,7 @@
 const { TEXT_MAP } = require('../../../ext/formats')
 const { CLIENT_PORT_KEY } = require('../../dd-trace/src/constants')
 const ProducerPlugin = require('../../dd-trace/src/plugins/producer')
-const { DsmPathwayCodec } = require('../../dd-trace/src/datastreams/pathway')
-const { getAmqpMessageSize } = require('../../dd-trace/src/datastreams/processor')
+const { DsmPathwayCodec, getAmqpMessageSize } = require('../../dd-trace/src/datastreams')
 const { getResourceName } = require('./util')
 
 class AmqplibProducerPlugin extends ProducerPlugin {
diff --git a/packages/datadog-plugin-avsc/src/schema_iterator.js b/packages/datadog-plugin-avsc/src/schema_iterator.js
@@ -10,7 +10,7 @@ const {
 const log = require('../../dd-trace/src/log')
 const {
   SchemaBuilder
-} = require('../../dd-trace/src/datastreams/schemas/schema_builder')
+} = require('../../dd-trace/src/datastreams')
 
 class SchemaExtractor {
   constructor (schema) {
diff --git a/packages/datadog-plugin-aws-sdk/src/services/kinesis.js b/packages/datadog-plugin-aws-sdk/src/services/kinesis.js
@@ -1,8 +1,5 @@
 'use strict'
-const {
-  getSizeOrZero
-} = require('../../../dd-trace/src/datastreams/processor')
-const { DsmPathwayCodec } = require('../../../dd-trace/src/datastreams/pathway')
+const { DsmPathwayCodec, getSizeOrZero } = require('../../../dd-trace/src/datastreams')
 const log = require('../../../dd-trace/src/log')
 const BaseAwsSdkPlugin = require('../base')
 const { storage } = require('../../../datadog-core')
diff --git a/packages/datadog-plugin-aws-sdk/src/services/sns.js b/packages/datadog-plugin-aws-sdk/src/services/sns.js
@@ -1,6 +1,5 @@
 'use strict'
-const { getHeadersSize } = require('../../../dd-trace/src/datastreams/processor')
-const { DsmPathwayCodec } = require('../../../dd-trace/src/datastreams/pathway')
+const { DsmPathwayCodec, getHeadersSize } = require('../../../dd-trace/src/datastreams')
 const log = require('../../../dd-trace/src/log')
 const BaseAwsSdkPlugin = require('../base')
 
diff --git a/packages/datadog-plugin-aws-sdk/src/services/sqs.js b/packages/datadog-plugin-aws-sdk/src/services/sqs.js
@@ -3,8 +3,7 @@
 const log = require('../../../dd-trace/src/log')
 const BaseAwsSdkPlugin = require('../base')
 const { storage } = require('../../../datadog-core')
-const { getHeadersSize } = require('../../../dd-trace/src/datastreams/processor')
-const { DsmPathwayCodec } = require('../../../dd-trace/src/datastreams/pathway')
+const { DsmPathwayCodec, getHeadersSize } = require('../../../dd-trace/src/datastreams')
 
 class Sqs extends BaseAwsSdkPlugin {
   static get id () { return 'sqs' }
diff --git a/packages/datadog-plugin-google-cloud-pubsub/src/consumer.js b/packages/datadog-plugin-google-cloud-pubsub/src/consumer.js
@@ -1,6 +1,6 @@
 'use strict'
 
-const { getMessageSize } = require('../../dd-trace/src/datastreams/processor')
+const { getMessageSize } = require('../../dd-trace/src/datastreams')
 const ConsumerPlugin = require('../../dd-trace/src/plugins/consumer')
 
 class GoogleCloudPubsubConsumerPlugin extends ConsumerPlugin {
diff --git a/packages/datadog-plugin-google-cloud-pubsub/src/producer.js b/packages/datadog-plugin-google-cloud-pubsub/src/producer.js
@@ -1,8 +1,7 @@
 'use strict'
 
 const ProducerPlugin = require('../../dd-trace/src/plugins/producer')
-const { DsmPathwayCodec } = require('../../dd-trace/src/datastreams/pathway')
-const { getHeadersSize } = require('../../dd-trace/src/datastreams/processor')
+const { DsmPathwayCodec, getHeadersSize } = require('../../dd-trace/src/datastreams')
 
 class GoogleCloudPubsubProducerPlugin extends ProducerPlugin {
   static get id () { return 'google-cloud-pubsub' }
diff --git a/packages/datadog-plugin-kafkajs/src/batch-consumer.js b/packages/datadog-plugin-kafkajs/src/batch-consumer.js
@@ -1,5 +1,5 @@
 const ConsumerPlugin = require('../../dd-trace/src/plugins/consumer')
-const { getMessageSize } = require('../../dd-trace/src/datastreams/processor')
+const { getMessageSize } = require('../../dd-trace/src/datastreams')
 
 class KafkajsBatchConsumerPlugin extends ConsumerPlugin {
   static get id () { return 'kafkajs' }
diff --git a/packages/datadog-plugin-kafkajs/src/consumer.js b/packages/datadog-plugin-kafkajs/src/consumer.js
@@ -1,7 +1,7 @@
 'use strict'
 
 const dc = require('dc-polyfill')
-const { getMessageSize } = require('../../dd-trace/src/datastreams/processor')
+const { getMessageSize } = require('../../dd-trace/src/datastreams')
 const ConsumerPlugin = require('../../dd-trace/src/plugins/consumer')
 
 const afterStartCh = dc.channel('dd-trace:kafkajs:consumer:afterStart')
diff --git a/packages/datadog-plugin-kafkajs/src/producer.js b/packages/datadog-plugin-kafkajs/src/producer.js
@@ -1,8 +1,7 @@
 'use strict'
 
 const ProducerPlugin = require('../../dd-trace/src/plugins/producer')
-const { DsmPathwayCodec } = require('../../dd-trace/src/datastreams/pathway')
-const { getMessageSize } = require('../../dd-trace/src/datastreams/processor')
+const { DsmPathwayCodec, getMessageSize } = require('../../dd-trace/src/datastreams')
 
 const BOOTSTRAP_SERVERS_KEY = 'messaging.kafka.bootstrap.servers'
 const MESSAGING_DESTINATION_KEY = 'messaging.destination.name'
diff --git a/packages/datadog-plugin-kafkajs/test/index.spec.js b/packages/datadog-plugin-kafkajs/test/index.spec.js
@@ -8,7 +8,7 @@ const { expectSomeSpan, withDefaults } = require('../../dd-trace/test/plugins/he
 const { ERROR_MESSAGE, ERROR_TYPE, ERROR_STACK } = require('../../dd-trace/src/constants')
 
 const { expectedSchema, rawExpectedSchema } = require('./naming')
-const DataStreamsContext = require('../../dd-trace/src/data_streams_context')
+const DataStreamsContext = require('../../dd-trace/src/datastreams/context')
 const { computePathwayHash } = require('../../dd-trace/src/datastreams/pathway')
 const { ENTRY_PARENT_HASH, DataStreamsProcessor } = require('../../dd-trace/src/datastreams/processor')
 
diff --git a/packages/datadog-plugin-protobufjs/src/schema_iterator.js b/packages/datadog-plugin-protobufjs/src/schema_iterator.js
@@ -10,7 +10,7 @@ const {
 const log = require('../../dd-trace/src/log')
 const {
   SchemaBuilder
-} = require('../../dd-trace/src/datastreams/schemas/schema_builder')
+} = require('../../dd-trace/src/datastreams')
 
 class SchemaExtractor {
   constructor (schema) {
diff --git a/packages/datadog-plugin-rhea/src/consumer.js b/packages/datadog-plugin-rhea/src/consumer.js
@@ -2,7 +2,7 @@
 
 const ConsumerPlugin = require('../../dd-trace/src/plugins/consumer')
 const { storage } = require('../../datadog-core')
-const { getAmqpMessageSize } = require('../../dd-trace/src/datastreams/processor')
+const { getAmqpMessageSize } = require('../../dd-trace/src/datastreams')
 
 class RheaConsumerPlugin extends ConsumerPlugin {
   static get id () { return 'rhea' }
diff --git a/packages/datadog-plugin-rhea/src/producer.js b/packages/datadog-plugin-rhea/src/producer.js
@@ -2,8 +2,7 @@
 
 const { CLIENT_PORT_KEY } = require('../../dd-trace/src/constants')
 const ProducerPlugin = require('../../dd-trace/src/plugins/producer')
-const { DsmPathwayCodec } = require('../../dd-trace/src/datastreams/pathway')
-const { getAmqpMessageSize } = require('../../dd-trace/src/datastreams/processor')
+const { getAmqpMessageSize, DsmPathwayCodec } = require('../../dd-trace/src/datastreams')
 
 class RheaProducerPlugin extends ProducerPlugin {
   static get id () { return 'rhea' }
diff --git a/packages/dd-trace/src/datastreams/checkpointer.js b/packages/dd-trace/src/datastreams/checkpointer.js
@@ -1,4 +1,4 @@
-const DataStreamsContext = require('./data_streams_context')
+const DataStreamsContext = require('./context')
 
 class DataStreamsCheckpointer {
   constructor (tracer) {
diff --git a/packages/dd-trace/src/datastreams/context.js b/packages/dd-trace/src/datastreams/context.js
@@ -1,5 +1,5 @@
-const { storage } = require('../../datadog-core')
-const log = require('./log')
+const { storage } = require('../../../datadog-core')
+const log = require('../log')
 
 function getDataStreamsContext () {
   const store = storage('legacy').getStore()
diff --git a/packages/dd-trace/src/datastreams/index.js b/packages/dd-trace/src/datastreams/index.js
@@ -0,0 +1,104 @@
+'use strict'
+
+const {
+  getAmqpMessageSize,
+  getHeadersSize,
+  getMessageSize,
+  getSizeOrZero
+} = require('./size')
+
+// This is only needed because DSM code is spread across existing tracing
+// plugins instead of having dedicated DSM plugins that are themselves
+// lazy loaded.
+//
+// TODO: Remove this when DSM has been moved to dedicaed plugins.
+function lazyClass (classGetter, methods = [], staticMethods = []) {
+  let constructorArgs
+  let ActiveClass
+
+  const LazyClass = function (...args) {
+    constructorArgs = args
+  }
+
+  const activate = () => {
+    return (ActiveClass = ActiveClass || classGetter())
+  }
+
+  for (const method of methods) {
+    LazyClass.prototype[method] = function (...args) {
+      const instance = activate() && new ActiveClass(...constructorArgs)
+
+      // Replace the whole prototype instead of only the method itself whenever
+      // any individual method is called to avoid running through this code
+      // again every time another method is called. This is not only more
+      // efficient but it also means that the class instance does not need to be
+      // stored for future calls to other methods.
+      Object.setPrototypeOf(this, instance)
+
+      return this[method](...args)
+    }
+  }
+
+  for (const method of staticMethods) {
+    LazyClass[method] = function (...args) {
+      LazyClass[method] = activate() && ActiveClass[method]
+
+      return LazyClass[method](...args)
+    }
+  }
+
+  return LazyClass
+}
+
+const DsmPathwayCodec = lazyClass(() => require('./pathway').DsmPathwayCodec, [], [
+  'encode',
+  'decode'
+])
+
+const DataStreamsCheckpointer = lazyClass(() => require('./checkpointer').DataStreamsCheckpointer, [
+  'setProduceCheckpoint',
+  'setConsumeCheckpoint'
+])
+
+const DataStreamsManager = lazyClass(() => require('./manager').DataStreamsManager, [
+  'setCheckpoint',
+  'decodeDataStreamsContext'
+])
+
+// TODO: Are all those methods actually public?
+const DataStreamsProcessor = lazyClass(() => require('./processor').DataStreamsProcessor, [
+  'onInterval',
+  'bucketFromTimestamp',
+  'recordCheckpoint',
+  'setCheckpoint',
+  'recordOffset',
+  'setOffset',
+  'setUrl',
+  'trySampleSchema',
+  'canSampleSchema',
+  'getSchema'
+])
+
+const SchemaBuilder = lazyClass(() => require('./schemas/schema_builder').SchemaBuilder, [
+  'build',
+  'addProperty',
+  'shouldExtractSchema'
+], [
+  'getCache',
+  'getSchemaDefinition',
+  'getSchema'
+])
+
+module.exports = {
+  DsmPathwayCodec,
+  DataStreamsCheckpointer,
+  DataStreamsManager,
+  DataStreamsProcessor,
+  SchemaBuilder,
+
+  // These are small functions so they are exposed directly and not lazy loaded.
+  getAmqpMessageSize,
+  getHeadersSize,
+  getMessageSize,
+  getSizeOrZero
+}
diff --git a/packages/dd-trace/src/datastreams/manager.js b/packages/dd-trace/src/datastreams/manager.js
@@ -0,0 +1,27 @@
+'use strict'
+
+const { DsmPathwayCodec } = require('./pathway')
+const DataStreamsContext = require('./context')
+
+class DataStreamsManager {
+  constructor (processor) {
+    this._dataStreamsProcessor = processor
+  }
+
+  setCheckpoint (edgeTags, span, payloadSize = 0) {
+    const ctx = this._dataStreamsProcessor.setCheckpoint(
+      edgeTags, span, DataStreamsContext.getDataStreamsContext(), payloadSize
+    )
+    DataStreamsContext.setDataStreamsContext(ctx)
+    return ctx
+  }
+
+  decodeDataStreamsContext (carrier) {
+    const ctx = DsmPathwayCodec.decode(carrier)
+    // we erase the previous context everytime we decode a new one
+    DataStreamsContext.setDataStreamsContext(ctx)
+    return ctx
+  }
+}
+
+module.exports = { DataStreamsManager }
diff --git a/packages/dd-trace/src/datastreams/processor.js b/packages/dd-trace/src/datastreams/processor.js
@@ -5,7 +5,7 @@ const { LogCollapsingLowestDenseDDSketch } = require('@datadog/sketches-js')
 const { DsmPathwayCodec } = require('./pathway')
 const { DataStreamsWriter } = require('./writer')
 const { computePathwayHash } = require('./pathway')
-const { types } = require('util')
+const { getAmqpMessageSize, getHeadersSize, getMessageSize, getSizeOrZero } = require('./size')
 const { PATHWAY_HASH } = require('../../../../ext/tags')
 const { SchemaBuilder } = require('./schemas/schema_builder')
 const { SchemaSampler } = require('./schemas/schema_sampler')
@@ -115,49 +115,6 @@ class StatsBucket {
   }
 }
 
-function getSizeOrZero (obj) {
-  if (typeof obj === 'string') {
-    return Buffer.from(obj, 'utf-8').length
-  }
-  if (types.isArrayBuffer(obj)) {
-    return obj.byteLength
-  }
-  if (Buffer.isBuffer(obj)) {
-    return obj.length
-  }
-  if (Array.isArray(obj) && obj.length > 0) {
-    if (typeof obj[0] === 'number') return Buffer.from(obj).length
-    let payloadSize = 0
-    obj.forEach(item => {
-      payloadSize += getSizeOrZero(item)
-    })
-    return payloadSize
-  }
-  if (obj !== null && typeof obj === 'object') {
-    try {
-      return getHeadersSize(obj)
-    } catch {
-      // pass
-    }
-  }
-  return 0
-}
-
-function getHeadersSize (headers) {
-  if (headers === undefined) return 0
-  return Object.entries(headers).reduce((prev, [key, val]) => getSizeOrZero(key) + getSizeOrZero(val) + prev, 0)
-}
-
-function getMessageSize (message) {
-  const { key, value, headers } = message
-  return getSizeOrZero(key) + getSizeOrZero(value) + getHeadersSize(headers)
-}
-
-function getAmqpMessageSize (message) {
-  const { headers, content } = message
-  return getSizeOrZero(content) + getHeadersSize(headers)
-}
-
 class TimeBuckets extends Map {
   forTime (time) {
     if (!this.has(time)) {
diff --git a/packages/dd-trace/src/datastreams/size.js b/packages/dd-trace/src/datastreams/size.js
@@ -0,0 +1,53 @@
+'use strict'
+
+const { types } = require('util')
+
+function getSizeOrZero (obj) {
+  if (typeof obj === 'string') {
+    return Buffer.from(obj, 'utf-8').length
+  }
+  if (types.isArrayBuffer(obj)) {
+    return obj.byteLength
+  }
+  if (Buffer.isBuffer(obj)) {
+    return obj.length
+  }
+  if (Array.isArray(obj) && obj.length > 0) {
+    if (typeof obj[0] === 'number') return Buffer.from(obj).length
+    let payloadSize = 0
+    obj.forEach(item => {
+      payloadSize += getSizeOrZero(item)
+    })
+    return payloadSize
+  }
+  if (obj !== null && typeof obj === 'object') {
+    try {
+      return getHeadersSize(obj)
+    } catch {
+      // pass
+    }
+  }
+  return 0
+}
+
+function getHeadersSize (headers) {
+  if (headers === undefined) return 0
+  return Object.entries(headers).reduce((prev, [key, val]) => getSizeOrZero(key) + getSizeOrZero(val) + prev, 0)
+}
+
+function getMessageSize (message) {
+  const { key, value, headers } = message
+  return getSizeOrZero(key) + getSizeOrZero(value) + getHeadersSize(headers)
+}
+
+function getAmqpMessageSize (message) {
+  const { headers, content } = message
+  return getSizeOrZero(content) + getHeadersSize(headers)
+}
+
+module.exports = {
+  getMessageSize,
+  getHeadersSize,
+  getSizeOrZero,
+  getAmqpMessageSize
+}
diff --git a/packages/dd-trace/src/opentracing/propagation/text_map_dsm.js b/packages/dd-trace/src/opentracing/propagation/text_map_dsm.js
diff --git a/packages/dd-trace/src/tracer.js b/packages/dd-trace/src/tracer.js
