From f38ee7776aeb846b2bb881e0c636806987b4a50d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Zugmeyer?= <benoit.zugmeyer@datadoghq.com>
Date: Wed, 12 Aug 2020 19:37:47 +0200
Subject: [PATCH] [RUMF-636] set initial trace id to the initial document
 resource

---
 packages/rum/src/performanceCollection.ts     | 23 +++++--
 .../rum/test/performanceCollection.spec.ts    | 67 +++++++++++++++++++
 2 files changed, 84 insertions(+), 6 deletions(-)

diff --git a/packages/rum/src/performanceCollection.ts b/packages/rum/src/performanceCollection.ts
index 45e7b31ae1..87b2628bc1 100644
--- a/packages/rum/src/performanceCollection.ts
+++ b/packages/rum/src/performanceCollection.ts
@@ -1,5 +1,6 @@
-import { DOM_EVENT, getRelativeTime, isNumber, monitor } from '@datadog/browser-core'
+import { DOM_EVENT, getRelativeTime, isNumber, monitor, ONE_MINUTE } from '@datadog/browser-core'
 
+import { getAPMDocumentData } from './getAPMDocumentData'
 import { LifeCycle, LifeCycleEventType } from './lifeCycle'
 import { FAKE_INITIAL_DOCUMENT } from './resourceUtils'
 
@@ -67,9 +68,8 @@ function supportPerformanceNavigationTimingEvent() {
 }
 
 export function startPerformanceCollection(lifeCycle: LifeCycle) {
-  retrieveInitialDocumentResourceTiming((timing) => {
-    handleRumPerformanceEntry(lifeCycle, timing)
-  })
+  handleRumPerformanceEntry(lifeCycle, retrieveInitialDocumentResourceTiming())
+
   if (supportPerformanceObject()) {
     handlePerformanceEntries(lifeCycle, performance.getEntries())
   }
@@ -99,9 +99,20 @@ export function startPerformanceCollection(lifeCycle: LifeCycle) {
   }
 }
 
-function retrieveInitialDocumentResourceTiming(callback: (timing: RumPerformanceResourceTiming) => void) {
+export const INITIAL_DOCUMENT_OUTDATED_TRACE_ID_THRESHOLD = 2 * ONE_MINUTE
+
+export function retrieveInitialDocumentResourceTiming() {
   let timing: RumPerformanceResourceTiming
+
+  let traceId
+  const apmDocumentData = getAPMDocumentData(document)
+  const now = Date.now()
+  if (apmDocumentData && apmDocumentData.traceTime > now - INITIAL_DOCUMENT_OUTDATED_TRACE_ID_THRESHOLD) {
+    traceId = apmDocumentData.traceId
+  }
+
   const forcedAttributes = {
+    traceId,
     entryType: 'resource' as const,
     initiatorType: FAKE_INITIAL_DOCUMENT,
   }
@@ -119,7 +130,7 @@ function retrieveInitialDocumentResourceTiming(callback: (timing: RumPerformance
       ...forcedAttributes,
     }
   }
-  callback(timing)
+  return timing
 }
 
 function retrieveNavigationTimingWhenLoaded(callback: (timing: RumPerformanceNavigationTiming) => void) {
diff --git a/packages/rum/test/performanceCollection.spec.ts b/packages/rum/test/performanceCollection.spec.ts
index cbeedbab04..d7c2789464 100644
--- a/packages/rum/test/performanceCollection.spec.ts
+++ b/packages/rum/test/performanceCollection.spec.ts
@@ -1,6 +1,10 @@
 import { isIE } from '@datadog/browser-core'
 
 import { restorePageVisibility, setPageVisibility } from '../../core/src/specHelper'
+import {
+  INITIAL_DOCUMENT_OUTDATED_TRACE_ID_THRESHOLD,
+  retrieveInitialDocumentResourceTiming,
+} from '../../rum/src/performanceCollection'
 import { setup, TestSetupBuilder } from './specHelper'
 
 describe('rum first_contentful_paint', () => {
@@ -34,3 +38,66 @@ describe('rum first_contentful_paint', () => {
     expect(stubBuilder.getEntryTypes()).toContain('paint')
   })
 })
+
+describe('rum initial document resource', () => {
+  let setupBuilder: TestSetupBuilder
+  beforeEach(() => {
+    setupBuilder = setup().withPerformanceCollection()
+  })
+
+  afterEach(() => {
+    setupBuilder.cleanup()
+  })
+
+  it('creates a resource timing for the initial document', () => {
+    const timing = retrieveInitialDocumentResourceTiming()
+    expect(timing.entryType).toBe('resource')
+    expect(timing.duration).toBeGreaterThan(0)
+  })
+
+  describe('initial trace id association', () => {
+    const nodesToRemove: Node[] = []
+
+    afterEach(() => {
+      nodesToRemove.forEach((node) => {
+        node.parentNode!.removeChild(node)
+      })
+      nodesToRemove.length = 0
+    })
+
+    it('should set a trace id found from an APM document comment', () => {
+      const comment = new Comment(`DATADOG;trace-id=123;trace-time=${Date.now()}`)
+      nodesToRemove.push(comment)
+      document.appendChild(comment)
+
+      const timing = retrieveInitialDocumentResourceTiming()
+      expect(timing.traceId).toBe('123')
+    })
+
+    it('should set a trace id found from meta tags', () => {
+      const metaTraceId = document.createElement('meta')
+      metaTraceId.name = 'dd-trace-id'
+      metaTraceId.content = '456'
+      document.head.appendChild(metaTraceId)
+      const metaTraceTime = document.createElement('meta')
+      metaTraceTime.name = 'dd-trace-time'
+      metaTraceTime.content = String(Date.now())
+      document.head.appendChild(metaTraceTime)
+      nodesToRemove.push(metaTraceId, metaTraceTime)
+
+      const timing = retrieveInitialDocumentResourceTiming()
+      expect(timing.traceId).toBe('456')
+    })
+
+    it('should ignore the trace id if it has been created too long ago', () => {
+      const comment = new Comment(
+        `DATADOG;trace-id=123;trace-time=${Date.now() - INITIAL_DOCUMENT_OUTDATED_TRACE_ID_THRESHOLD}`
+      )
+      nodesToRemove.push(comment)
+      document.appendChild(comment)
+
+      const timing = retrieveInitialDocumentResourceTiming()
+      expect(timing.traceId).toBe(undefined)
+    })
+  })
+})