add manual checkpoint parameter to consume DSM API (#6190)

michael-zhao459 · web-flow · commit 1f496ad19b32 · 2025-07-29T13:17:16.000-04:00
* made manual checkpoint configurable

* add comments

* lint

* rewrote tests
diff --git a/packages/dd-trace/src/datastreams/checkpointer.js b/packages/dd-trace/src/datastreams/checkpointer.js
@@ -9,6 +9,12 @@ class DataStreamsCheckpointer {
     this.dsmProcessor = tracer._dataStreamsProcessor
   }
 
+  /**
+   * @param {string} type - The type of the checkpoint, usually the streaming technology being used.
+   *                       Examples include kafka, kinesis, sns etc.
+   * @param {string} target - The target of data. This can be a topic, exchange or stream name.
+   * @param {Object} carrier - The carrier object to inject context into.
+   */
   setProduceCheckpoint (type, target, carrier) {
     if (!this.config.dsmEnabled) return
 
@@ -23,14 +29,29 @@ class DataStreamsCheckpointer {
     this.tracer.inject(ctx, 'text_map_dsm', carrier)
   }
 
-  setConsumeCheckpoint (type, source, carrier) {
+  /**
+   * @param {string} type - The type of the checkpoint, usually the streaming technology being used.
+   *                       Examples include kafka, kinesis, sns etc.
+   * @param {string} source - The source of data. This can be a topic, exchange or stream name.
+   * @param {Object} carrier - The carrier object to extract context from.
+   * @param {boolean} [manualCheckpoint=true] - Whether this checkpoint was manually set. Keep true if manually
+   *                                           instrumenting. Manual instrumentation always overrides automatic
+   *                                           instrumentation in the case a call is both manually and automatically
+   *                                           instrumented.
+   */
+  setConsumeCheckpoint (type, source, carrier, manualCheckpoint = true) {
     if (!this.config.dsmEnabled) return
 
     const parentCtx = this.tracer.extract('text_map_dsm', carrier)
     DataStreamsContext.setDataStreamsContext(parentCtx)
 
+    const tags = ['type:' + type, 'topic:' + source, 'direction:in']
+    if (manualCheckpoint) {
+      tags.push('manual_checkpoint:true')
+    }
+
     const ctx = this.dsmProcessor.setCheckpoint(
-      ['type:' + type, 'topic:' + source, 'direction:in', 'manual_checkpoint:true'],
+      tags,
       null,
       parentCtx,
       null
diff --git a/packages/dd-trace/test/datastreams/data_streams_checkpointer.spec.js b/packages/dd-trace/test/datastreams/data_streams_checkpointer.spec.js
@@ -70,4 +70,26 @@ describe('data streams checkpointer manual api', () => {
 
     expect(DSM_CONTEXT_HEADER in headers).to.equal(true)
   })
+
+  it('should set manual checkpoint when setConsumeCheckpoint is called without additional parameters', function () {
+    const headers = {}
+    const mockSetCheckpoint = sinon.stub().returns({ hash: Buffer.from([1, 2, 3, 4]) })
+
+    tracer._tracer._dataStreamsProcessor.setCheckpoint = mockSetCheckpoint
+
+    tracer.dataStreamsCheckpointer.setConsumeCheckpoint('kinesis', 'stream-123', headers)
+    const calledTags = mockSetCheckpoint.getCall(0).args[0]
+    expect(calledTags).to.include('manual_checkpoint:true')
+  })
+
+  it('should set an automatic checkpoint when setConsumeCheckpoint is called with manualCheckpoint:false', function () {
+    const headers = {}
+    const mockSetCheckpoint = sinon.stub().returns({ hash: Buffer.from([1, 2, 3, 4]) })
+
+    tracer._tracer._dataStreamsProcessor.setCheckpoint = mockSetCheckpoint
+
+    tracer.dataStreamsCheckpointer.setConsumeCheckpoint('kinesis', 'stream-123', headers, false)
+    const calledTags = mockSetCheckpoint.getCall(0).args[0]
+    expect(calledTags).to.not.include('manual_checkpoint:true')
+  })
 })
