/
index.ts
112 lines (100 loc) · 3.32 KB
/
index.ts
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
import {
getLogger,
getNamespace,
isAutomaticMode,
middleware,
Segment,
setSegment,
utils,
} from 'aws-xray-sdk-core';
import { ServerResponse } from 'http';
import { Context, Middleware, Next } from 'koa';
const { IncomingRequestData } = middleware;
/**
* Creates the koa middleware function
*
* @export
* @param {string} defaultName - Service name for XRay
* @returns {Middleware} - The middleware function
*/
export default function createMiddleware(defaultName: string): Middleware {
if (!defaultName || typeof defaultName !== 'string')
throw new Error(
'Default segment name was not supplied. Please provide a string.'
);
middleware.setDefaultName(defaultName);
return async function xrayMiddleware(ctx: Context, next: Next) {
try {
const amznTraceHeader = middleware.processHeaders(ctx);
const name = middleware.resolveName(ctx.host);
ctx.segment = new Segment(
name,
amznTraceHeader.Root,
amznTraceHeader.Parent
);
middleware.resolveSampling(
amznTraceHeader,
ctx.segment,
ctx as unknown as ServerResponse
);
ctx.segment.addIncomingRequestData(new IncomingRequestData(ctx.req));
getLogger().debug(
`Starting koa segment: { url: ${ctx.url}, name: ${
ctx.segment.name
}, trace_id: ${ctx.segment.trace_id}, id: ${
ctx.segment.id
}, sampled: ${!ctx.segment.notTraced} }`
);
if (isAutomaticMode()) {
const ns = getNamespace();
ns.bindEmitter(ctx.req);
ns.bindEmitter(ctx.res);
ns.runAndReturn(async function namespace() {
setSegment(ctx.segment);
await next();
if (ctx.segment) {
ctx.segment.close();
getLogger().debug(
`Closed koa segment successfully: { url: ${ctx.url}, name: ${
ctx.segment.name
}, trace_id: ${ctx.segment.trace_id}, id: ${
ctx.segment.id
}, sampled: ${!ctx.segment.notTraced} }`
);
if (ctx.status === 429) ctx.segment.addThrottleFlag();
if (utils.getCauseTypeFromHttpStatus(ctx.status))
ctx.segment[
utils.getCauseTypeFromHttpStatus(ctx.status) as string
] = true;
ctx.segment.http.close(ctx);
ctx.segment.close();
}
});
} else {
await next();
if (ctx.status === 429) ctx.segment.addThrottleFlag();
if (utils.getCauseTypeFromHttpStatus(ctx.status))
ctx.segment[utils.getCauseTypeFromHttpStatus(ctx.status) as string] =
true;
ctx.segment.http.close(ctx);
ctx.segment.close();
}
} catch (err) {
if (ctx.segment && err) {
if (ctx.status === 429) ctx.segment.addThrottleFlag();
if (utils.getCauseTypeFromHttpStatus(ctx.status))
ctx.segment[utils.getCauseTypeFromHttpStatus(ctx.status) as string] =
true;
ctx.segment.http.close(ctx);
ctx.segment.close(err);
getLogger().debug(
`Closed koa segment with error: { url: ${ctx.url}, name: ${
ctx.segment.name
}, trace_id: ${ctx.segment.trace_id}, id: ${
ctx.segment.id
}, sampled: ${!ctx.segment.notTraced} }`
);
}
}
};
}