Skip to content

Commit 991e77a

Browse files
rochdevrochdev
authored
fix shimmer double wrapping for non-method functions (#5329)
1 parent 8068772 commit 991e77a

1 file changed

Lines changed: 95 additions & 95 deletions

File tree

‎packages/datadog-shimmer/src/shimmer.js‎

Lines changed: 95 additions & 95 deletions
Original file line numberDiff line numberDiff line change
@@ -26,10 +26,10 @@ function copyProperties (original, wrapped) {
2626

2727
function wrapFunction (original, wrapper) {
2828
if (typeof original === 'function') assertNotClass(original)
29-
// TODO This needs to be re-done so that this and wrapMethod are distinct.
30-
const target = { func: original }
31-
wrapMethod(target, 'func', wrapper, typeof original !== 'function')
32-
let delegate = target.func
29+
30+
let delegate = safeMode
31+
? safeWrapper(original, wrapper)
32+
: wrapper(original)
3333

3434
const shim = function shim () {
3535
return delegate.apply(this, arguments)
@@ -85,98 +85,10 @@ function wrapMethod (target, name, wrapper, noAssert) {
8585
}
8686

8787
const original = target[name]
88-
let wrapped
89-
90-
if (safeMode && original) {
91-
// In this mode, we make a best-effort attempt to handle errors that are thrown
92-
// by us, rather than wrapped code. With such errors, we log them, and then attempt
93-
// to return the result as if no wrapping was done at all.
94-
//
95-
// Caveats:
96-
// * If the original function is called in a later iteration of the event loop,
97-
// and we throw _then_, then it won't be caught by this. In practice, we always call
98-
// the original function synchronously, so this is not a problem.
99-
// * While async errors are dealt with here, errors in callbacks are not. This
100-
// is because we don't necessarily know _for sure_ that any function arguments
101-
// are wrapped by us. We could wrap them all anyway and just make that assumption,
102-
// or just assume that the last argument is always a callback set by us if it's a
103-
// function, but those don't seem like things we can rely on. We could add a
104-
// `shimmer.markCallbackAsWrapped()` function that's a no-op outside safe-mode,
105-
// but that means modifying every instrumentation. Even then, the complexity of
106-
// this code increases because then we'd need to effectively do the reverse of
107-
// what we're doing for synchronous functions. This is a TODO.
108-
109-
// We're going to hold on to current callState in this variable in this scope,
110-
// which is fine because any time we reference it, we're referencing it synchronously.
111-
// We'll use it in the our wrapper (which, again, is called syncrhonously), and in the
112-
// errorHandler, which will already have been bound to this callState.
113-
let currentCallState
114-
115-
// Rather than calling the original function directly from the shim wrapper, we wrap
116-
// it again so that we can track if it was called and if it returned. This is because
117-
// we need to know if an error was thrown by the original function, or by us.
118-
// We could do this inside the `wrapper` function defined below, which would simplify
119-
// managing the callState, but then we'd be calling `wrapper` on each invocation, so
120-
// instead we do it here, once.
121-
const innerWrapped = wrapper(function (...args) {
122-
// We need to stash the callState here because of recursion.
123-
const callState = currentCallState
124-
callState.startCall()
125-
const retVal = original.apply(this, args)
126-
if (isPromise(retVal)) {
127-
retVal.then(callState.endCall.bind(callState))
128-
} else {
129-
callState.endCall(retVal)
130-
}
131-
return retVal
132-
})
133-
134-
// This is the crux of what we're doing in safe mode. It handles errors
135-
// that _we_ cause, by logging them, and transparently providing results
136-
// as if no wrapping was done at all. That means detecting (via callState)
137-
// whether the function has already run or not, and if it has, returning
138-
// the result, and otherwise calling the original function unwrapped.
139-
const handleError = function (args, callState, e) {
140-
if (callState.completed) {
141-
// error was thrown after original function returned/resolved, so
142-
// it was us. log it.
143-
log.error('Shimmer error was thrown after original function returned/resolved', e)
144-
// original ran and returned something. return it.
145-
return callState.retVal
146-
}
147-
148-
if (!callState.called) {
149-
// error was thrown before original function was called, so
150-
// it was us. log it.
151-
log.error('Shimmer error was thrown before original function was called', e)
152-
// original never ran. call it unwrapped.
153-
return original.apply(this, args)
154-
}
155-
156-
// error was thrown during original function execution, so
157-
// it was them. throw.
158-
throw e
159-
}
88+
const wrapped = safeMode && original
89+
? safeWrapper(original, wrapper)
90+
: wrapper(original)
16091

161-
// The wrapped function is the one that will be called by the user.
162-
// It calls our version of the original function, which manages the
163-
// callState. That way when we use the errorHandler, it can tell where
164-
// the error originated.
165-
wrapped = function (...args) {
166-
currentCallState = new CallState()
167-
const errorHandler = handleError.bind(this, args, currentCallState)
168-
169-
try {
170-
const retVal = innerWrapped.apply(this, args)
171-
return isPromise(retVal) ? retVal.catch(errorHandler) : retVal
172-
} catch (e) {
173-
return errorHandler(e)
174-
}
175-
}
176-
} else {
177-
// In non-safe mode, we just wrap the original function directly.
178-
wrapped = wrapper(original)
179-
}
18092
const descriptor = Object.getOwnPropertyDescriptor(target, name)
18193

18294
const attributes = {
@@ -212,6 +124,94 @@ function wrapMethod (target, name, wrapper, noAssert) {
212124
return target
213125
}
214126

127+
function safeWrapper (original, wrapper) {
128+
// In this mode, we make a best-effort attempt to handle errors that are thrown
129+
// by us, rather than wrapped code. With such errors, we log them, and then attempt
130+
// to return the result as if no wrapping was done at all.
131+
//
132+
// Caveats:
133+
// * If the original function is called in a later iteration of the event loop,
134+
// and we throw _then_, then it won't be caught by this. In practice, we always call
135+
// the original function synchronously, so this is not a problem.
136+
// * While async errors are dealt with here, errors in callbacks are not. This
137+
// is because we don't necessarily know _for sure_ that any function arguments
138+
// are wrapped by us. We could wrap them all anyway and just make that assumption,
139+
// or just assume that the last argument is always a callback set by us if it's a
140+
// function, but those don't seem like things we can rely on. We could add a
141+
// `shimmer.markCallbackAsWrapped()` function that's a no-op outside safe-mode,
142+
// but that means modifying every instrumentation. Even then, the complexity of
143+
// this code increases because then we'd need to effectively do the reverse of
144+
// what we're doing for synchronous functions. This is a TODO.
145+
146+
// We're going to hold on to current callState in this variable in this scope,
147+
// which is fine because any time we reference it, we're referencing it synchronously.
148+
// We'll use it in the our wrapper (which, again, is called syncrhonously), and in the
149+
// errorHandler, which will already have been bound to this callState.
150+
let currentCallState
151+
152+
// Rather than calling the original function directly from the shim wrapper, we wrap
153+
// it again so that we can track if it was called and if it returned. This is because
154+
// we need to know if an error was thrown by the original function, or by us.
155+
// We could do this inside the `wrapper` function defined below, which would simplify
156+
// managing the callState, but then we'd be calling `wrapper` on each invocation, so
157+
// instead we do it here, once.
158+
const innerWrapped = wrapper(function (...args) {
159+
// We need to stash the callState here because of recursion.
160+
const callState = currentCallState
161+
callState.startCall()
162+
const retVal = original.apply(this, args)
163+
if (isPromise(retVal)) {
164+
retVal.then(callState.endCall.bind(callState))
165+
} else {
166+
callState.endCall(retVal)
167+
}
168+
return retVal
169+
})
170+
171+
// This is the crux of what we're doing in safe mode. It handles errors
172+
// that _we_ cause, by logging them, and transparently providing results
173+
// as if no wrapping was done at all. That means detecting (via callState)
174+
// whether the function has already run or not, and if it has, returning
175+
// the result, and otherwise calling the original function unwrapped.
176+
const handleError = function (args, callState, e) {
177+
if (callState.completed) {
178+
// error was thrown after original function returned/resolved, so
179+
// it was us. log it.
180+
log.error('Shimmer error was thrown after original function returned/resolved', e)
181+
// original ran and returned something. return it.
182+
return callState.retVal
183+
}
184+
185+
if (!callState.called) {
186+
// error was thrown before original function was called, so
187+
// it was us. log it.
188+
log.error('Shimmer error was thrown before original function was called', e)
189+
// original never ran. call it unwrapped.
190+
return original.apply(this, args)
191+
}
192+
193+
// error was thrown during original function execution, so
194+
// it was them. throw.
195+
throw e
196+
}
197+
198+
// The wrapped function is the one that will be called by the user.
199+
// It calls our version of the original function, which manages the
200+
// callState. That way when we use the errorHandler, it can tell where
201+
// the error originated.
202+
return function (...args) {
203+
currentCallState = new CallState()
204+
const errorHandler = handleError.bind(this, args, currentCallState)
205+
206+
try {
207+
const retVal = innerWrapped.apply(this, args)
208+
return isPromise(retVal) ? retVal.catch(errorHandler) : retVal
209+
} catch (e) {
210+
return errorHandler(e)
211+
}
212+
}
213+
}
214+
215215
function wrap (target, name, wrapper) {
216216
return typeof name === 'function'
217217
? wrapFn(target, name)

0 commit comments

Comments
 (0)