From 8b6c11a365a7409168e94ac9b26b5f2f5c18e3d2 Mon Sep 17 00:00:00 2001
From: Alan Storm <astorm@users.noreply.github.com>
Date: Thu, 12 May 2022 09:52:44 -0700
Subject: [PATCH] feat: setRecorded for traceparent/tracecontext (#2687)

* feat: setRecorded for traceparent/tracecontext

Part of https://github.com/elastic/apm-agent-nodejs/issues/2302
---
 lib/tracecontext/index.js             |  4 ++++
 lib/tracecontext/traceparent.js       | 14 ++++++++++++--
 test/tracecontext/traceparent.test.js | 14 ++++++++++++++
 3 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/lib/tracecontext/index.js b/lib/tracecontext/index.js
index 83274ae49e..f37ef21a56 100644
--- a/lib/tracecontext/index.js
+++ b/lib/tracecontext/index.js
@@ -82,6 +82,10 @@ class TraceContext {
       setter(carrier, 'tracestate', tracestate)
     }
   }
+
+  setRecorded () {
+    return this.traceparent.setRecorded()
+  }
 }
 
 TraceContext.FLAGS = TraceParent.FLAGS
diff --git a/lib/tracecontext/traceparent.js b/lib/tracecontext/traceparent.js
index 8a9cf86227..e38258a8a0 100644
--- a/lib/tracecontext/traceparent.js
+++ b/lib/tracecontext/traceparent.js
@@ -107,8 +107,10 @@ class TraceParent {
   constructor (buffer) {
     this[bufferSymbol] = buffer
     Object.defineProperty(this, 'recorded', {
-      value: !!(buffer[OFFSETS.flags] & FLAGS.recorded),
-      enumerable: true
+      enumerable: true,
+      get: function () {
+        return !!(this[bufferSymbol][OFFSETS.flags] & FLAGS.recorded)
+      }
     })
 
     defineLazyProp(this, 'version', hexSliceFn(buffer, OFFSETS.version, OFFSETS.traceId))
@@ -144,6 +146,14 @@ class TraceParent {
     return makeChild(Buffer.from(this[bufferSymbol]))
   }
 
+  setRecorded (value) {
+    if (value) {
+      this[bufferSymbol][OFFSETS.flags] |= FLAGS.recorded
+    } else {
+      this[bufferSymbol][OFFSETS.flags] &= ~FLAGS.recorded
+    }
+  }
+
   toString () {
     return `${this.version}-${this.traceId}-${this.id}-${this.flags}`
   }
diff --git a/test/tracecontext/traceparent.test.js b/test/tracecontext/traceparent.test.js
index c25b21f3c5..77ce275ff7 100644
--- a/test/tracecontext/traceparent.test.js
+++ b/test/tracecontext/traceparent.test.js
@@ -176,3 +176,17 @@ test('ensureParentId', t => {
 
   t.end()
 })
+
+test('setRecorded', t => {
+  const traceParent = TraceParent.fromString(header)
+
+  t.ok(traceParent.recorded)
+
+  traceParent.setRecorded(false)
+  t.ok(!traceParent.recorded)
+
+  traceParent.setRecorded(true)
+  t.ok(traceParent.recorded)
+
+  t.end()
+})