Skip to content

Commit 5c8a3a2

Browse files
ardatanardatan
committed
feat(otel): ability to propagate otel context
1 parent c7c5f5e commit 5c8a3a2

File tree

2 files changed

+84
-38
lines changed

2 files changed

+84
-38
lines changed

‎.changeset/yellow-countries-mate.md‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
'@envelop/opentelemetry': patch
3+
---
4+
5+
Ability to propagate otel context

‎packages/plugins/opentelemetry/src/index.ts‎

Lines changed: 79 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,24 @@ type PluginContext = {
4444
[tracingSpanSymbol]: opentelemetry.Span;
4545
};
4646

47+
export const otelContextMap = new WeakMap<any, opentelemetry.Context>();
48+
49+
export function getCurrentOtelContext(graphqlContext: any): opentelemetry.Context {
50+
let otelContext = otelContextMap.get(graphqlContext);
51+
52+
if (!otelContext) {
53+
otelContext = opentelemetry.context.active();
54+
otelContextMap.set(graphqlContext, otelContext);
55+
}
56+
57+
return otelContext;
58+
}
59+
60+
export function setCurrentOtelContext(graphqlContext: any, otelContext: opentelemetry.Context) {
61+
otelContextMap.set(graphqlContext, otelContext);
62+
return otelContext;
63+
}
64+
4765
export const useOpenTelemetry = (
4866
options: TracingOptions,
4967
tracingProvider?: TracerProvider,
@@ -70,7 +88,7 @@ export const useOpenTelemetry = (
7088
useOnResolve(({ info, context, args }) => {
7189
const parentSpan = spanByContext.get(context);
7290
if (parentSpan) {
73-
const ctx = opentelemetry.trace.setSpan(opentelemetry.context.active(), parentSpan);
91+
const ctx = opentelemetry.trace.setSpan(getCurrentOtelContext(context), parentSpan);
7492
const { fieldName, returnType, parentType } = info;
7593

7694
const resolverSpan = tracer.startSpan(
@@ -103,7 +121,12 @@ export const useOpenTelemetry = (
103121
);
104122
}
105123
},
106-
onExecute({ args }) {
124+
onExecute({ args, executeFn, setExecuteFn }) {
125+
setExecuteFn(function wrappedExecuteFnWithOtelCtx(args) {
126+
return opentelemetry.context.with(getCurrentOtelContext(args.contextValue), () =>
127+
executeFn(args),
128+
);
129+
});
107130
const operationAst = getOperationAST(args.document, args.operationName);
108131
if (!operationAst) {
109132
return;
@@ -122,24 +145,29 @@ export const useOpenTelemetry = (
122145
isDocumentLoggable = false;
123146
}
124147
const operationName = operationAst.name?.value || 'anonymous';
125-
const executionSpan = tracer.startSpan(`${spanPrefix}${operationType}.${operationName}`, {
126-
kind: spanKind,
127-
attributes: {
128-
...spanAdditionalAttributes,
129-
[AttributeName.EXECUTION_OPERATION_NAME]: operationName,
130-
[AttributeName.EXECUTION_OPERATION_TYPE]: operationType,
131-
[AttributeName.EXECUTION_OPERATION_DOCUMENT]: isDocumentLoggable
132-
? getDocumentString(args.document, print)
133-
: undefined,
134-
...(options.variables
135-
? { [AttributeName.EXECUTION_VARIABLES]: JSON.stringify(args.variableValues ?? {}) }
136-
: {}),
148+
const currOtelContext = getCurrentOtelContext(args.contextValue);
149+
const executionSpan = tracer.startSpan(
150+
`${spanPrefix}${operationType}.${operationName}`,
151+
{
152+
kind: spanKind,
153+
attributes: {
154+
...spanAdditionalAttributes,
155+
[AttributeName.EXECUTION_OPERATION_NAME]: operationName,
156+
[AttributeName.EXECUTION_OPERATION_TYPE]: operationType,
157+
[AttributeName.EXECUTION_OPERATION_DOCUMENT]: isDocumentLoggable
158+
? getDocumentString(args.document, print)
159+
: undefined,
160+
...(options.variables
161+
? { [AttributeName.EXECUTION_VARIABLES]: JSON.stringify(args.variableValues ?? {}) }
162+
: {}),
163+
},
137164
},
138-
});
165+
currOtelContext,
166+
);
139167

140-
const otelContext = opentelemetry.trace.setSpan(
141-
opentelemetry.context.active(),
142-
executionSpan,
168+
setCurrentOtelContext(
169+
args.contextValue,
170+
opentelemetry.trace.setSpan(currOtelContext, executionSpan),
143171
);
144172

145173
const resultCbs: OnExecuteHookResult<PluginContext> = {
@@ -149,7 +177,8 @@ export const useOpenTelemetry = (
149177
executionSpan.setAttribute(AttributeName.EXECUTION_RESULT, JSON.stringify(result));
150178
}
151179
if (options.traceIdInResult) {
152-
setResult(addTraceIdToResult(otelContext, result, options.traceIdInResult));
180+
const currOtelContext = getCurrentOtelContext(args.contextValue);
181+
setResult(addTraceIdToResult(currOtelContext, result, options.traceIdInResult));
153182
}
154183
markError(executionSpan, result);
155184
executionSpan.end();
@@ -159,7 +188,8 @@ export const useOpenTelemetry = (
159188
// handles async iterator
160189
onNext: ({ result, setResult }) => {
161190
if (options.traceIdInResult) {
162-
setResult(addTraceIdToResult(otelContext, result, options.traceIdInResult));
191+
const currOtelContext = getCurrentOtelContext(args.contextValue);
192+
setResult(addTraceIdToResult(currOtelContext, result, options.traceIdInResult));
163193
}
164194
markError(executionSpan, result);
165195
},
@@ -176,7 +206,12 @@ export const useOpenTelemetry = (
176206

177207
return resultCbs;
178208
},
179-
onSubscribe({ args }) {
209+
onSubscribe({ args, subscribeFn, setSubscribeFn }) {
210+
setSubscribeFn(function wrappedSubscribeFnWithOtelCtx(args) {
211+
return opentelemetry.context.with(getCurrentOtelContext(args.contextValue), () =>
212+
subscribeFn(args),
213+
);
214+
});
180215
const operationAst = getOperationAST(args.document, args.operationName);
181216
if (!operationAst) {
182217
return;
@@ -190,25 +225,30 @@ export const useOpenTelemetry = (
190225
} else {
191226
isDocumentLoggable = false;
192227
}
228+
const currOtelContext = getCurrentOtelContext(args.contextValue);
193229
const operationName = operationAst.name?.value || 'anonymous';
194-
const subscriptionSpan = tracer.startSpan(`${operationType}.${operationName}`, {
195-
kind: spanKind,
196-
attributes: {
197-
...spanAdditionalAttributes,
198-
[AttributeName.EXECUTION_OPERATION_NAME]: operationName,
199-
[AttributeName.EXECUTION_OPERATION_TYPE]: operationType,
200-
[AttributeName.EXECUTION_OPERATION_DOCUMENT]: isDocumentLoggable
201-
? getDocumentString(args.document, print)
202-
: undefined,
203-
...(options.variables
204-
? { [AttributeName.EXECUTION_VARIABLES]: JSON.stringify(args.variableValues ?? {}) }
205-
: {}),
230+
const subscriptionSpan = tracer.startSpan(
231+
`${operationType}.${operationName}`,
232+
{
233+
kind: spanKind,
234+
attributes: {
235+
...spanAdditionalAttributes,
236+
[AttributeName.EXECUTION_OPERATION_NAME]: operationName,
237+
[AttributeName.EXECUTION_OPERATION_TYPE]: operationType,
238+
[AttributeName.EXECUTION_OPERATION_DOCUMENT]: isDocumentLoggable
239+
? getDocumentString(args.document, print)
240+
: undefined,
241+
...(options.variables
242+
? { [AttributeName.EXECUTION_VARIABLES]: JSON.stringify(args.variableValues ?? {}) }
243+
: {}),
244+
},
206245
},
207-
});
246+
currOtelContext,
247+
);
208248

209-
const otelContext = opentelemetry.trace.setSpan(
210-
opentelemetry.context.active(),
211-
subscriptionSpan,
249+
setCurrentOtelContext(
250+
args.contextValue,
251+
opentelemetry.trace.setSpan(currOtelContext, subscriptionSpan),
212252
);
213253

214254
const resultCbs: OnSubscribeHookResult<PluginContext> = {
@@ -220,7 +260,8 @@ export const useOpenTelemetry = (
220260
// handles async iterator
221261
onNext: ({ result, setResult }) => {
222262
if (options.traceIdInResult) {
223-
setResult(addTraceIdToResult(otelContext, result, options.traceIdInResult));
263+
const currOtelContext = getCurrentOtelContext(args.contextValue);
264+
setResult(addTraceIdToResult(currOtelContext, result, options.traceIdInResult));
224265
}
225266
markError(subscriptionSpan, result);
226267
},

0 commit comments

Comments
 (0)