[RUMF-824] support Request instances in tracing (DataDog#684)

WonderInventions · Jan 29, 2021 · 3b3c12a · 3b3c12a
1 parent e6861b8
commit 3b3c12a
Show file tree

Hide file tree

Showing 6 changed files with 108 additions and 44 deletions.
diff --git a/packages/core/src/browser/fetchProxy.ts b/packages/core/src/browser/fetchProxy.ts
@@ -14,6 +14,7 @@ export interface FetchProxy<
 export interface FetchStartContext {
   method: string
   startTime: number
+  input: RequestInfo
   init?: RequestInit
   url: string
 
@@ -77,6 +78,7 @@ function proxyFetch() {
 
     const context: FetchStartContext = {
       init,
+      input,
       method,
       startTime,
       url,
@@ -106,7 +108,7 @@ function proxyFetch() {
     }
     beforeSendCallbacks.forEach((callback) => callback(context))
 
-    const responsePromise = originalFetch.call(this, input, context.init)
+    const responsePromise = originalFetch.call(this, context.input, context.init)
     responsePromise.then(monitor(reportFetch), monitor(reportFetch))
     return responsePromise
   })

diff --git a/packages/core/src/tools/specHelper.ts b/packages/core/src/tools/specHelper.ts
@@ -61,8 +61,8 @@ export function stubFetch(): FetchStubManager {
       resolve = res
       reject = rej
     })
-    ;(promise as FetchStubPromise).resolveWith = async (response: ResponseStub) => {
-      const resolved = resolve({
+    ;(promise as FetchStubPromise).resolveWith = (response: ResponseStub) => {
+      resolve({
         ...response,
         clone: () => {
           const cloned = {
@@ -75,14 +75,12 @@ export function stubFetch(): FetchStubManager {
           }
           return cloned as Response
         },
-      }) as Promise<ResponseStub>
+      })
       onRequestEnd()
-      return resolved
     }
-    ;(promise as FetchStubPromise).rejectWith = async (error: Error) => {
-      const rejected = reject(error) as Promise<Error>
+    ;(promise as FetchStubPromise).rejectWith = (error: Error) => {
+      reject(error)
       onRequestEnd()
-      return rejected
     }
     return promise
   }) as typeof window.fetch
@@ -106,8 +104,8 @@ export interface ResponseStub extends Partial<Response> {
 export type FetchStub = (input: RequestInfo, init?: RequestInit) => FetchStubPromise
 
 export interface FetchStubPromise extends Promise<Response> {
-  resolveWith: (response: ResponseStub) => Promise<ResponseStub>
-  rejectWith: (error: Error) => Promise<Error>
+  resolveWith: (response: ResponseStub) => void
+  rejectWith: (error: Error) => void
 }
 
 class StubXhr {

diff --git a/packages/rum-core/src/domain/tracing/tracer.spec.ts b/packages/rum-core/src/domain/tracing/tracer.spec.ts
@@ -189,6 +189,48 @@ describe('tracer', () => {
       ])
     })
 
+    it('should preserve original headers contained in a Request instance', () => {
+      const request = new Request(document.location.origin, {
+        headers: {
+          foo: 'bar',
+        },
+      })
+
+      const context: Partial<RumFetchCompleteContext> = {
+        ...ALLOWED_DOMAIN_CONTEXT,
+        input: request,
+      }
+
+      const tracer = startTracer(configuration as Configuration)
+      tracer.traceFetch(context)
+
+      expect(context.init).toBe(undefined)
+      expect(context.input).not.toBe(request)
+      expect(headersAsArray((context.input as Request).headers)).toEqual([
+        ['foo', 'bar'],
+        ...tracingHeadersAsArrayFor(context.traceId!, context.spanId!),
+      ])
+      expect(headersAsArray(request.headers)).toEqual([['foo', 'bar']])
+    })
+
+    it('should ignore headers from a Request instance if other headers are set', () => {
+      const context: Partial<RumFetchCompleteContext> = {
+        ...ALLOWED_DOMAIN_CONTEXT,
+        init: { headers: { 'x-init-header': 'baz' } },
+        input: new Request(document.location.origin, {
+          headers: { 'x-request-header': 'bar' },
+        }),
+      }
+
+      const tracer = startTracer(configuration as Configuration)
+      tracer.traceFetch(context)
+
+      expect(context.init!.headers).toEqual([
+        ['x-init-header', 'baz'],
+        ...tracingHeadersAsArrayFor(context.traceId!, context.spanId!),
+      ])
+    })
+
     it('should not trace request on disallowed domain', () => {
       const context: Partial<RumFetchCompleteContext> = { ...DISALLOWED_DOMAIN_CONTEXT }
 
@@ -277,5 +319,13 @@ function tracingHeadersFor(traceId: TraceIdentifier, spanId: TraceIdentifier) {
 }
 
 function tracingHeadersAsArrayFor(traceId: TraceIdentifier, spanId: TraceIdentifier) {
-  return objectEntries(tracingHeadersFor(traceId, spanId)) as string[][]
+  return objectEntries(tracingHeadersFor(traceId, spanId)) as Array<[string, string]>
+}
+
+function headersAsArray(headers: Headers) {
+  const result: Array<[string, string]> = []
+  headers.forEach((value, key) => {
+    result.push([key, value])
+  })
+  return result
 }
diff --git a/packages/rum-core/src/domain/tracing/tracer.ts b/packages/rum-core/src/domain/tracing/tracer.ts
@@ -28,22 +28,29 @@ export function startTracer(configuration: Configuration): Tracer {
     clearTracingIfCancelled,
     traceFetch: (context) =>
       injectHeadersIfTracingAllowed(configuration, context, (tracingHeaders: TracingHeaders) => {
-        context.init = { ...context.init }
-        const headers: string[][] = []
-        if (context.init.headers instanceof Headers) {
-          context.init.headers.forEach((value, key) => {
-            headers.push([key, value])
-          })
-        } else if (Array.isArray(context.init.headers)) {
-          context.init.headers.forEach((header) => {
-            headers.push(header)
-          })
-        } else if (context.init.headers) {
-          Object.keys(context.init.headers).forEach((key) => {
-            headers.push([key, (context.init!.headers as Record<string, string>)[key]])
+        if (context.input instanceof Request && !context.init?.headers) {
+          context.input = new Request(context.input)
+          Object.keys(tracingHeaders).forEach((key) => {
+            ;(context.input as Request).headers.append(key, tracingHeaders[key])
           })
+        } else {
+          context.init = { ...context.init }
+          const headers: string[][] = []
+          if (context.init.headers instanceof Headers) {
+            context.init.headers.forEach((value, key) => {
+              headers.push([key, value])
+            })
+          } else if (Array.isArray(context.init.headers)) {
+            context.init.headers.forEach((header) => {
+              headers.push(header)
+            })
+          } else if (context.init.headers) {
+            Object.keys(context.init.headers).forEach((key) => {
+              headers.push([key, (context.init!.headers as Record<string, string>)[key]])
+            })
+          }
+          context.init.headers = headers.concat(objectEntries(tracingHeaders) as string[][])
         }
-        context.init.headers = headers.concat(objectEntries(tracingHeaders) as string[][])
       }),
     traceXhr: (context, xhr) =>
       injectHeadersIfTracingAllowed(configuration, context, (tracingHeaders: TracingHeaders) => {

diff --git a/test/e2e/lib/helpers/browser.ts b/test/e2e/lib/helpers/browser.ts
@@ -71,17 +71,3 @@ export async function sendXhr(url: string, headers: string[][] = []): Promise<st
   }
   return result.response
 }
-
-export async function sendFetch(url: string, headers: string[][] = []): Promise<string> {
-  return browserExecuteAsync(
-    // tslint:disable-next-line: no-shadowed-variable
-    (url, headers, done) => {
-      window
-        .fetch(url, { headers })
-        .then((response) => response.text())
-        .then(done)
-    },
-    url,
-    headers
-  )
-}
diff --git a/test/e2e/scenario/rum/tracing.scenario.ts b/test/e2e/scenario/rum/tracing.scenario.ts
@@ -1,5 +1,5 @@
 import { createTest, EventRegistry } from '../../lib/framework'
-import { sendFetch, sendXhr } from '../../lib/helpers/browser'
+import { browserExecuteAsync, sendXhr } from '../../lib/helpers/browser'
 import { flushEvents } from '../../lib/helpers/sdk'
 
 describe('tracing', () => {
@@ -18,10 +18,31 @@ describe('tracing', () => {
   createTest('trace fetch')
     .withRum({ service: 'Service', allowedTracingOrigins: ['LOCATION_ORIGIN'] })
     .run(async ({ events }) => {
-      const rawHeaders = await sendFetch(`/headers`, [
-        ['x-foo', 'bar'],
-        ['x-foo', 'baz'],
-      ])
+      const rawHeaders = await browserExecuteAsync<string>((done) => {
+        window
+          .fetch('/headers', {
+            headers: [
+              ['x-foo', 'bar'],
+              ['x-foo', 'baz'],
+            ],
+          })
+          .then((response) => response.text())
+          .then(done)
+      })
+      checkRequestHeaders(rawHeaders)
+      await flushEvents()
+      await checkTraceAssociatedToRumEvent(events)
+    })
+
+  createTest('trace fetch with Request argument')
+    .withRum({ service: 'Service', allowedTracingOrigins: ['LOCATION_ORIGIN'] })
+    .run(async ({ events }) => {
+      const rawHeaders = await browserExecuteAsync<string>((done) => {
+        window
+          .fetch(new Request('/headers', { headers: { 'x-foo': 'bar, baz' } }))
+          .then((response) => response.text())
+          .then(done)
+      })
       checkRequestHeaders(rawHeaders)
       await flushEvents()
       await checkTraceAssociatedToRumEvent(events)