-
Notifications
You must be signed in to change notification settings - Fork 8.1k
/
index.ts
116 lines (92 loc) · 3.02 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
113
114
115
116
/*
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
* or more contributor license agreements. Licensed under the Elastic License
* 2.0 and the Server Side Public License, v 1; you may not use this file except
* in compliance with, at your election, the Elastic License 2.0 or the Server
* Side Public License, v 1.
*/
import agent from 'elastic-apm-node';
import asyncHooks from 'async_hooks';
export interface SpanOptions {
name: string;
type?: string;
subtype?: string;
labels?: Record<string, string>;
intercept?: boolean;
}
type Span = Exclude<typeof agent.currentSpan, undefined | null>;
export function parseSpanOptions(optionsOrName: SpanOptions | string) {
const options = typeof optionsOrName === 'string' ? { name: optionsOrName } : optionsOrName;
return options;
}
const runInNewContext = <T extends (...args: any[]) => any>(cb: T): ReturnType<T> => {
const resource = new asyncHooks.AsyncResource('fake_async');
return resource.runInAsyncScope(cb);
};
export async function withSpan<T>(
optionsOrName: SpanOptions | string,
cb: (span?: Span) => Promise<T>
): Promise<T> {
const options = parseSpanOptions(optionsOrName);
const { name, type, subtype, labels, intercept } = options;
if (!agent.isStarted()) {
return cb();
}
let createdSpan: Span | undefined;
// When a span starts, it's marked as the active span in its context.
// When it ends, it's not untracked, which means that if a span
// starts directly after this one ends, the newly started span is a
// child of this span, even though it should be a sibling.
// To mitigate this, we queue a microtask by awaiting a promise.
if (!intercept) {
await Promise.resolve();
createdSpan = agent.startSpan(name) ?? undefined;
if (!createdSpan) {
return cb();
}
}
// If a span is created in the same context as the span that we just
// started, it will be a sibling, not a child. E.g., the Elasticsearch span
// that is created when calling search() happens in the same context. To
// mitigate this we create a new context.
return runInNewContext(() => {
const promise = cb(createdSpan);
let span: Span | undefined = createdSpan;
if (intercept) {
span = agent.currentSpan ?? undefined;
}
if (!span) {
return promise;
}
const targetedSpan = span;
if (name) {
targetedSpan.name = name;
}
// @ts-ignore
if (type) {
targetedSpan.type = type;
}
if (subtype) {
targetedSpan.subtype = subtype;
}
if (labels) {
targetedSpan.addLabels(labels);
}
return promise
.then((res) => {
if (!targetedSpan.outcome || targetedSpan.outcome === 'unknown') {
targetedSpan.outcome = 'success';
}
return res;
})
.catch((err) => {
if (!targetedSpan.outcome || targetedSpan.outcome === 'unknown') {
targetedSpan.outcome = 'failure';
}
throw err;
})
.finally(() => {
targetedSpan.end();
});
});
}