Skip to content

Commit 6c9facb

Browse files
DarcyRaynerDDDarcyRaynerDD
authored
Merge branch 'master' into darcy.rayner/add-node-12-support
2 parents 076b2a1 + 49c98ce commit 6c9facb

File tree

10 files changed

+248
-15
lines changed

10 files changed

+248
-15
lines changed

‎src/index.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ export const defaultConfig: Config = {
4141
debugLogging: false,
4242
enhancedMetrics: false,
4343
logForwarding: false,
44+
mergeDatadogXrayTraces: false,
4445
shouldRetryMetrics: false,
4546
siteURL: "",
4647
} as const;

‎src/trace/constants.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,10 @@ export enum SampleMode {
44
AUTO_KEEP = 1,
55
USER_KEEP = 2,
66
}
7+
export enum Source {
8+
Xray = "xray",
9+
Event = "event",
10+
}
711

812
export const traceIDHeader = "x-datadog-trace-id";
913
export const parentIDHeader = "x-datadog-parent-id";

‎src/trace/context.spec.ts‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { LogLevel, setLogLevel } from "../utils";
2-
import { SampleMode, xrayBaggageSubsegmentKey, xraySubsegmentNamespace } from "./constants";
2+
import { SampleMode, xrayBaggageSubsegmentKey, xraySubsegmentNamespace, Source } from "./constants";
33
import {
44
convertToAPMParentID,
55
convertToAPMTraceID,
@@ -105,6 +105,7 @@ describe("convertTraceContext", () => {
105105
parentID: "797643193680388254",
106106
sampleMode: SampleMode.USER_KEEP,
107107
traceID: "4110911582297405557",
108+
source: Source.Xray,
108109
});
109110
});
110111
it("returns undefined if traceID is invalid", () => {
@@ -137,9 +138,10 @@ describe("readTraceContextFromXray", () => {
137138
parentID: "797643193680388254",
138139
sampleMode: SampleMode.USER_KEEP,
139140
traceID: "4110911582297405557",
141+
source: Source.Xray,
140142
});
141143
});
142-
it("will parse a trace context from the xray, with sampling turned off", () => {
144+
it("will ignore a trace context from the xray, when sampling is turned off", () => {
143145
currentSegment = {
144146
id: "0b11cc4230d3e09e",
145147
notTraced: true,
@@ -151,6 +153,7 @@ describe("readTraceContextFromXray", () => {
151153
parentID: "797643193680388254",
152154
sampleMode: SampleMode.USER_REJECT,
153155
traceID: "4110911582297405557",
156+
source: Source.Xray,
154157
});
155158
});
156159
it("returns undefined when trace header isn't in environment", () => {
@@ -172,6 +175,7 @@ describe("readTraceFromEvent", () => {
172175
parentID: "797643193680388254",
173176
sampleMode: SampleMode.USER_KEEP,
174177
traceID: "4110911582297405557",
178+
source: Source.Event,
175179
});
176180
});
177181
it("can read well formed headers with mixed casing", () => {
@@ -186,6 +190,7 @@ describe("readTraceFromEvent", () => {
186190
parentID: "797643193680388254",
187191
sampleMode: SampleMode.USER_KEEP,
188192
traceID: "4110911582297405557",
193+
source: Source.Event,
189194
});
190195
});
191196
it("returns undefined when missing trace id", () => {
@@ -373,6 +378,7 @@ describe("extractTraceContext", () => {
373378
parentID: "797643193680388251",
374379
sampleMode: SampleMode.USER_KEEP,
375380
traceID: "4110911582297405551",
381+
source: Source.Event,
376382
});
377383
});
378384
it("returns trace read from env if no headers present", () => {
@@ -386,6 +392,7 @@ describe("extractTraceContext", () => {
386392
parentID: "797643193680388254",
387393
sampleMode: SampleMode.USER_KEEP,
388394
traceID: "4110911582297405557",
395+
source: "xray",
389396
});
390397
});
391398
it("returns trace read from env if no headers present", () => {
@@ -399,6 +406,7 @@ describe("extractTraceContext", () => {
399406
parentID: "797643193680388254",
400407
sampleMode: SampleMode.USER_KEEP,
401408
traceID: "4110911582297405557",
409+
source: "xray",
402410
});
403411
});
404412

‎src/trace/context.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import {
66
parentIDHeader,
77
SampleMode,
88
samplingPriorityHeader,
9+
Source,
910
traceIDHeader,
1011
xrayBaggageSubsegmentKey,
1112
xraySubsegmentKey,
@@ -23,6 +24,7 @@ export interface TraceContext {
2324
traceID: string;
2425
parentID: string;
2526
sampleMode: SampleMode;
27+
source: Source;
2628
}
2729

2830
export interface StepFunctionContext {
@@ -110,6 +112,7 @@ export function readTraceFromEvent(event: any): TraceContext | undefined {
110112
return {
111113
parentID,
112114
sampleMode,
115+
source: Source.Event,
113116
traceID,
114117
};
115118
}
@@ -189,6 +192,7 @@ export function convertTraceContext(traceHeader: XRayTraceHeader): TraceContext
189192
return {
190193
parentID,
191194
sampleMode,
195+
source: Source.Xray,
192196
traceID,
193197
};
194198
}

‎src/trace/dd-trace-utils.spec.ts‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
describe("isTracerInitialized", () => {
2+
let isTracerInitialized: () => boolean;
3+
let tracer: any;
4+
beforeEach(() => {
5+
isTracerInitialized = require("./dd-trace-utils").isTracerInitialized;
6+
tracer = require("dd-trace");
7+
process.env["AWS_LAMBDA_FUNCTION_NAME"] = "my-lambda";
8+
});
9+
afterEach(() => {
10+
jest.resetModules();
11+
delete process.env["AWS_LAMBDA_FUNCTION_NAME"];
12+
});
13+
it("should return true when tracer has been initialised", () => {
14+
tracer.init();
15+
expect(isTracerInitialized()).toBeTruthy();
16+
});
17+
it("should return false when tracer hasn't been initialised", () => {
18+
expect(isTracerInitialized()).toBeFalsy();
19+
});
20+
});

‎src/trace/dd-trace-utils.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
import Tracer from "dd-trace";
2+
3+
export function isTracerInitialized() {
4+
// TODO, waiting for a better way from APM to tell whether tracer has been initialised.
5+
const tracer = Tracer as any;
6+
return tracer !== undefined && tracer._tracer !== undefined && "_service" in tracer._tracer;
7+
}

‎src/trace/listener.spec.ts‎

Lines changed: 159 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,159 @@
1+
import { TraceListener } from "./listener";
2+
import { Source } from "./constants";
3+
4+
let mockWrap: jest.Mock<any, any>;
5+
let mockExtract: jest.Mock<any, any>;
6+
let mockTraceHeaders: Record<string, string> | undefined = undefined;
7+
let mockTraceSource: Source | undefined = undefined;
8+
9+
jest.mock("dd-trace", () => {
10+
mockWrap = jest.fn().mockImplementation((name, options, func) => func);
11+
mockExtract = jest.fn().mockImplementation((_, val) => val);
12+
return {
13+
wrap: mockWrap,
14+
extract: mockExtract,
15+
};
16+
});
17+
18+
jest.mock("./trace-context-service", () => {
19+
class MockTraceContextService {
20+
get traceSource() {
21+
return mockTraceSource;
22+
}
23+
get currentTraceHeaders() {
24+
return mockTraceHeaders;
25+
}
26+
}
27+
return {
28+
TraceContextService: MockTraceContextService,
29+
};
30+
});
31+
32+
describe("TraceListener", () => {
33+
const defaultConfig = { autoPatchHTTP: true, mergeDatadogXrayTraces: false };
34+
const context = {
35+
invokedFunctionArn: "arn:aws:lambda:us-east-1:123456789101:function:my-lambda",
36+
awsRequestId: "1234",
37+
functionName: "my-lambda",
38+
};
39+
beforeEach(() => {
40+
mockWrap.mockClear();
41+
mockExtract.mockClear();
42+
mockTraceHeaders = undefined;
43+
mockTraceSource = undefined;
44+
});
45+
46+
it("wraps dd-trace span around invocation", async () => {
47+
const listener = new TraceListener(defaultConfig, "handler.my-handler");
48+
listener.onStartInvocation({}, context as any);
49+
const unwrappedFunc = () => {};
50+
const wrappedFunc = listener.onWrap(unwrappedFunc);
51+
wrappedFunc();
52+
await listener.onCompleteInvocation();
53+
54+
expect(mockWrap).toHaveBeenCalledWith(
55+
"aws.lambda",
56+
{
57+
resource: "handler.my-handler",
58+
tags: {
59+
cold_start: true,
60+
function_arn: "arn:aws:lambda:us-east-1:123456789101:function:my-lambda",
61+
request_id: "1234",
62+
resource_names: "my-lambda",
63+
},
64+
},
65+
unwrappedFunc,
66+
);
67+
});
68+
69+
it("wraps dd-trace span around invocation, with trace context from event", async () => {
70+
const listener = new TraceListener(defaultConfig, "handler.my-handler");
71+
mockTraceHeaders = {
72+
"x-datadog-parent-id": "797643193680388251",
73+
"x-datadog-sampling-priority": "2",
74+
"x-datadog-trace-id": "4110911582297405551",
75+
};
76+
mockTraceSource = Source.Event;
77+
listener.onStartInvocation({}, context as any);
78+
const unwrappedFunc = () => {};
79+
const wrappedFunc = listener.onWrap(unwrappedFunc);
80+
wrappedFunc();
81+
await listener.onCompleteInvocation();
82+
83+
expect(mockWrap).toHaveBeenCalledWith(
84+
"aws.lambda",
85+
{
86+
resource: "handler.my-handler",
87+
tags: {
88+
cold_start: true,
89+
function_arn: "arn:aws:lambda:us-east-1:123456789101:function:my-lambda",
90+
request_id: "1234",
91+
resource_names: "my-lambda",
92+
},
93+
childOf: mockTraceHeaders,
94+
},
95+
unwrappedFunc,
96+
);
97+
});
98+
99+
it("wraps dd-trace span around invocation, without trace context from xray", async () => {
100+
const listener = new TraceListener(defaultConfig, "handler.my-handler");
101+
mockTraceHeaders = {
102+
"x-datadog-parent-id": "797643193680388251",
103+
"x-datadog-sampling-priority": "2",
104+
"x-datadog-trace-id": "4110911582297405551",
105+
};
106+
mockTraceSource = Source.Xray;
107+
108+
listener.onStartInvocation({}, context as any);
109+
const unwrappedFunc = () => {};
110+
const wrappedFunc = listener.onWrap(unwrappedFunc);
111+
wrappedFunc();
112+
await listener.onCompleteInvocation();
113+
114+
expect(mockWrap).toHaveBeenCalledWith(
115+
"aws.lambda",
116+
{
117+
resource: "handler.my-handler",
118+
tags: {
119+
cold_start: true,
120+
function_arn: "arn:aws:lambda:us-east-1:123456789101:function:my-lambda",
121+
request_id: "1234",
122+
resource_names: "my-lambda",
123+
},
124+
},
125+
unwrappedFunc,
126+
);
127+
});
128+
129+
it("wraps dd-trace span around invocation, with trace context from xray when mergeDatadogXrayTraces is enabled", async () => {
130+
const listener = new TraceListener({ ...defaultConfig, mergeDatadogXrayTraces: true }, "handler.my-handler");
131+
mockTraceHeaders = {
132+
"x-datadog-parent-id": "797643193680388251",
133+
"x-datadog-sampling-priority": "2",
134+
"x-datadog-trace-id": "4110911582297405551",
135+
};
136+
mockTraceSource = Source.Xray;
137+
138+
listener.onStartInvocation({}, context as any);
139+
const unwrappedFunc = () => {};
140+
const wrappedFunc = listener.onWrap(unwrappedFunc);
141+
wrappedFunc();
142+
await listener.onCompleteInvocation();
143+
144+
expect(mockWrap).toHaveBeenCalledWith(
145+
"aws.lambda",
146+
{
147+
resource: "handler.my-handler",
148+
tags: {
149+
cold_start: true,
150+
function_arn: "arn:aws:lambda:us-east-1:123456789101:function:my-lambda",
151+
request_id: "1234",
152+
resource_names: "my-lambda",
153+
},
154+
childOf: mockTraceHeaders,
155+
},
156+
unwrappedFunc,
157+
);
158+
});
159+
});

0 commit comments

Comments
 (0)