add dsm for google pub sub

packages/datadog-plugin-google-cloud-pubsub/src/consumer.js

@@ -1,5 +1,6 @@
 'use strict'
 
+const { getMessageSize } = require('../../dd-trace/src/datastreams/processor')
 const ConsumerPlugin = require('../../dd-trace/src/plugins/consumer')
 
 class GoogleCloudPubsubConsumerPlugin extends ConsumerPlugin {
@@ -11,7 +12,7 @@ class GoogleCloudPubsubConsumerPlugin extends ConsumerPlugin {
     const topic = subscription.metadata && subscription.metadata.topic
     const childOf = this.tracer.extract('text_map', message.attributes) || null
 
-    this.startSpan({
+    const span = this.startSpan({
       childOf,
       resource: topic,
       type: 'worker',
@@ -23,6 +24,12 @@ class GoogleCloudPubsubConsumerPlugin extends ConsumerPlugin {
         'pubsub.ack': 0
       }
     })
+    if (this.config.dsmEnabled && message?.attributes) {
+      const payloadSize = getMessageSize(message)
+      this.tracer.decodeDataStreamsContext(message.attributes)
+      this.tracer
+        .setCheckpoint(['direction:in', `topic:${topic}`, 'type:google-pubsub'], span, payloadSize)
+    }
   }
 
   finish (message) {

packages/datadog-plugin-google-cloud-pubsub/src/producer.js

@@ -1,6 +1,8 @@
 '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')
 
 class GoogleCloudPubsubProducerPlugin extends ProducerPlugin {
   static get id () { return 'google-cloud-pubsub' }
@@ -25,6 +27,12 @@ class GoogleCloudPubsubProducerPlugin extends ProducerPlugin {
         msg.attributes = {}
       }
       this.tracer.inject(span, 'text_map', msg.attributes)
+      if (this.config.dsmEnabled) {
+        const payloadSize = getHeadersSize(msg)
+        const dataStreamsContext = this.tracer
+          .setCheckpoint(['direction:out', `topic:${topic}`, 'type:google-pubsub'], span, payloadSize)
+        DsmPathwayCodec.encode(dataStreamsContext, msg.attributes)
+      }
     }
   }
 }

packages/datadog-plugin-google-cloud-pubsub/test/index.spec.js

@@ -6,10 +6,28 @@ const id = require('../../dd-trace/src/id')
 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 { computePathwayHash } = require('../../dd-trace/src/datastreams/pathway')
+const { ENTRY_PARENT_HASH, DataStreamsProcessor } = require('../../dd-trace/src/datastreams/processor')
 
 // The roundtrip to the pubsub emulator takes time. Sometimes a *long* time.
 const TIMEOUT = 30000
 
+const dsmFullTopic = 'projects/test-project-42764574591eb894/topics/dsm-topic'
+const dsmTopicName = 'dsm-topic'
+const expectedProducerHash = computePathwayHash(
+  'test',
+  'tester',
+  ['direction:out', 'topic:' + dsmFullTopic, 'type:google-pubsub'],
+  ENTRY_PARENT_HASH
+)
+const expectedConsumerHash = computePathwayHash(
+  'test',
+  'tester',
+  ['direction:in', 'topic:' + dsmFullTopic, 'type:google-pubsub'],
+  expectedProducerHash
+)
+
 describe('Plugin', () => {
   let tracer
 
@@ -18,6 +36,7 @@ describe('Plugin', () => {
 
     before(() => {
       process.env.PUBSUB_EMULATOR_HOST = 'localhost:8081'
+      process.env.DD_DATA_STREAMS_ENABLED = true
     })
 
     after(() => {
@@ -37,7 +56,7 @@ describe('Plugin', () => {
 
       describe('without configuration', () => {
         beforeEach(() => {
-          return agent.load('google-cloud-pubsub')
+          return agent.load('google-cloud-pubsub', { dsmEnabled: false })
         })
 
         beforeEach(() => {
@@ -296,7 +315,8 @@ describe('Plugin', () => {
       describe('with configuration', () => {
         beforeEach(() => {
           return agent.load('google-cloud-pubsub', {
-            service: 'a_test_service'
+            service: 'a_test_service',
+            dsmEnabled: false
           })
         })
 
@@ -322,6 +342,98 @@ describe('Plugin', () => {
         })
       })
 
+      describe('data stream monitoring', () => {
+        let dsmTopic
+        let sub
+        let consume
+
+        beforeEach(() => {
+          return agent.load('google-cloud-pubsub', {
+            dsmEnabled: true
+          })
+        })
+
+        beforeEach(async () => {
+          const { PubSub } = require(`../../../versions/@google-cloud/pubsub@${version}`).get()
+          project = getProjectId()
+          resource = `projects/${project}/topics/${dsmTopicName}`
+          pubsub = new PubSub({ projectId: project })
+          tracer.use('google-cloud-pubsub', { dsmEnabled: true })
+
+          dsmTopic = await pubsub.createTopic(dsmTopicName)
+          dsmTopic = dsmTopic[0]
+          sub = await dsmTopic.createSubscription('DSM')
+          sub = sub[0]
+          consume = function (cb) {
+            sub.on('message', cb)
+          }
+        })
+
+        describe('should set a DSM checkpoint', () => {
+          let setDataStreamsContextSpy
+
+          beforeEach(() => {
+            setDataStreamsContextSpy = sinon.spy(DataStreamsContext, 'setDataStreamsContext')
+          })
+
+          afterEach(() => {
+            setDataStreamsContextSpy.restore()
+          })
+
+          it('on produce', async () => {
+            await publish(dsmTopic, { data: Buffer.from('DSM produce checkpoint') })
+
+            let testPassed
+            for (const checkpointCall of setDataStreamsContextSpy.args) {
+              if (checkpointCall[0].hash === expectedProducerHash) {
+                testPassed = true
+                break
+              }
+            }
+            expect(testPassed).to.equal(true)
+          })
+
+          it('on consume', async () => {
+            await publish(dsmTopic, { data: Buffer.from('DSM consume checkpoint') })
+            await consume(async () => {
+              let testPassed
+              for (const checkpointCall of setDataStreamsContextSpy.args) {
+                if (checkpointCall[0].hash === expectedConsumerHash) {
+                  testPassed = true
+                  break
+                }
+              }
+              expect(testPassed).to.equal(true)
+            })
+          })
+        })
+
+        describe('it should set a message payload size', () => {
+          let recordCheckpointSpy
+
+          beforeEach(() => {
+            recordCheckpointSpy = sinon.spy(DataStreamsProcessor.prototype, 'recordCheckpoint')
+          })
+
+          afterEach(() => {
+            DataStreamsProcessor.prototype.recordCheckpoint.restore()
+          })
+
+          it('when producing a message', async () => {
+            await publish(dsmTopic, { data: Buffer.from('DSM produce payload size') })
+            expect(recordCheckpointSpy.args[0][0].hasOwnProperty('payloadSize'))
+          })
+
+          it('when consuming a message', async () => {
+            await publish(dsmTopic, { data: Buffer.from('DSM consume payload size') })
+
+            await consume(async () => {
+              expect(recordCheckpointSpy.args[0][0].hasOwnProperty('payloadSize'))
+            })
+          })
+        })
+      })
+
       function expectSpanWithDefaults (expected) {
         const prefixedResource = [expected.meta['pubsub.method'], resource].filter(x => x).join(' ')
         const service = expected.meta['pubsub.method'] ? 'test-pubsub' : 'test'