Skip to content

Commit

Permalink
This is very scary
Browse files Browse the repository at this point in the history
  • Loading branch information
lforst committed Sep 20, 2024
1 parent dbfd8e2 commit b7786ac
Show file tree
Hide file tree
Showing 7 changed files with 561 additions and 107 deletions.
6 changes: 6 additions & 0 deletions packages/vercel-edge/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,13 @@
"access": "public"
},
"dependencies": {
"@opentelemetry/api": "^1.9.0",
"@opentelemetry/core": "^1.25.1",
"@opentelemetry/resources": "^1.26.0",
"@opentelemetry/sdk-trace-base": "^1.26.0",
"@opentelemetry/semantic-conventions": "^1.27.0",
"@sentry/core": "8.30.0",
"@sentry/opentelemetry": "8.30.0",
"@sentry/types": "8.30.0",
"@sentry/utils": "8.30.0"
},
Expand Down
324 changes: 324 additions & 0 deletions packages/vercel-edge/src/async-local-storage-context-manager.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,324 @@
/*
* Copyright The OpenTelemetry Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

// Taken from:
// - https://github.com/open-telemetry/opentelemetry-js/blob/6515ed8098333646a63a74a8c0150cc2daf520db/packages/opentelemetry-context-async-hooks/src/AbstractAsyncHooksContextManager.ts
// - https://github.com/open-telemetry/opentelemetry-js/blob/6515ed8098333646a63a74a8c0150cc2daf520db/packages/opentelemetry-context-async-hooks/src/AsyncLocalStorageContextManager.ts

import { EventEmitter } from 'events';
import { ROOT_CONTEXT } from '@opentelemetry/api';
import type { Context, ContextManager } from '@opentelemetry/api';
import { logger } from '@sentry/utils';
import { DEBUG_BUILD } from './debug-build';

interface AsyncLocalStorage<T> {
getStore(): T | undefined;
// eslint-disable-next-line @typescript-eslint/no-explicit-any
run<R, TArgs extends any[]>(store: T, callback: (...args: TArgs) => R, ...args: TArgs): R;
disable(): void;
}

type Func<T> = (...args: unknown[]) => T;

/**
* Store a map for each event of all original listeners and their "patched"
* version. So when a listener is removed by the user, the corresponding
* patched function will be also removed.
*/
interface PatchMap {
[name: string]: WeakMap<Func<void>, Func<void>>;
}

const ADD_LISTENER_METHODS = [
'addListener' as const,
'on' as const,
'once' as const,
'prependListener' as const,
'prependOnceListener' as const,
];

abstract class AbstractAsyncHooksContextManager implements ContextManager {
private readonly _kOtListeners = Symbol('OtListeners');
private _wrapped = false;

/**
* Binds a the certain context or the active one to the target function and then returns the target
* @param context A context (span) to be bind to target
* @param target a function or event emitter. When target or one of its callbacks is called,
* the provided context will be used as the active context for the duration of the call.
*/
public bind<T>(context: Context, target: T): T {
if (target instanceof EventEmitter) {
return this._bindEventEmitter(context, target);
}

if (typeof target === 'function') {
// @ts-expect-error This is vendored
return this._bindFunction(context, target);
}
return target;
}

/**
* By default, EventEmitter call their callback with their context, which we do
* not want, instead we will bind a specific context to all callbacks that
* go through it.
* @param context the context we want to bind
* @param ee EventEmitter an instance of EventEmitter to patch
*/
private _bindEventEmitter<T extends EventEmitter>(context: Context, ee: T): T {
const map = this._getPatchMap(ee);
if (map !== undefined) return ee;
this._createPatchMap(ee);

// patch methods that add a listener to propagate context
ADD_LISTENER_METHODS.forEach(methodName => {
if (ee[methodName] === undefined) return;
ee[methodName] = this._patchAddListener(ee, ee[methodName], context);
});
// patch methods that remove a listener
if (typeof ee.removeListener === 'function') {
// eslint-disable-next-line @typescript-eslint/unbound-method
ee.removeListener = this._patchRemoveListener(ee, ee.removeListener);
}
if (typeof ee.off === 'function') {
// eslint-disable-next-line @typescript-eslint/unbound-method
ee.off = this._patchRemoveListener(ee, ee.off);
}
// patch method that remove all listeners
if (typeof ee.removeAllListeners === 'function') {
// eslint-disable-next-line @typescript-eslint/unbound-method
ee.removeAllListeners = this._patchRemoveAllListeners(ee, ee.removeAllListeners);
}
return ee;
}

/**
* Patch methods that remove a given listener so that we match the "patched"
* version of that listener (the one that propagate context).
* @param ee EventEmitter instance
* @param original reference to the patched method
*/
// eslint-disable-next-line @typescript-eslint/no-explicit-any
private _patchRemoveListener(ee: EventEmitter, original: (...args: any[]) => any): any {
// eslint-disable-next-line @typescript-eslint/no-this-alias
const contextManager = this;
return function (this: never, event: string, listener: Func<void>) {
const events = contextManager._getPatchMap(ee)?.[event];
if (events === undefined) {
return original.call(this, event, listener);
}
const patchedListener = events.get(listener);
return original.call(this, event, patchedListener || listener);
};
}

/**
* Patch methods that remove all listeners so we remove our
* internal references for a given event.
* @param ee EventEmitter instance
* @param original reference to the patched method
*/
// eslint-disable-next-line @typescript-eslint/no-explicit-any
private _patchRemoveAllListeners(ee: EventEmitter, original: (...args: any[]) => any): any {
// eslint-disable-next-line @typescript-eslint/no-this-alias
const contextManager = this;
return function (this: never, event: string) {
const map = contextManager._getPatchMap(ee);
if (map !== undefined) {
if (arguments.length === 0) {
contextManager._createPatchMap(ee);
} else if (map[event] !== undefined) {
// eslint-disable-next-line @typescript-eslint/no-dynamic-delete
delete map[event];
}
}
// eslint-disable-next-line prefer-rest-params
return original.apply(this, arguments);
};
}

/**
* Patch methods on an event emitter instance that can add listeners so we
* can force them to propagate a given context.
* @param ee EventEmitter instance
* @param original reference to the patched method
* @param [context] context to propagate when calling listeners
*/
// eslint-disable-next-line @typescript-eslint/no-explicit-any
private _patchAddListener(ee: EventEmitter, original: (...args: any[]) => any, context: Context): any {
// eslint-disable-next-line @typescript-eslint/no-this-alias
const contextManager = this;
return function (this: never, event: string, listener: Func<void>) {
/**
* This check is required to prevent double-wrapping the listener.
* The implementation for ee.once wraps the listener and calls ee.on.
* Without this check, we would wrap that wrapped listener.
* This causes an issue because ee.removeListener depends on the onceWrapper
* to properly remove the listener. If we wrap their wrapper, we break
* that detection.
*/
if (contextManager._wrapped) {
return original.call(this, event, listener);
}
let map = contextManager._getPatchMap(ee);
if (map === undefined) {
map = contextManager._createPatchMap(ee);
}
let listeners = map[event];
if (listeners === undefined) {
listeners = new WeakMap();
map[event] = listeners;
}
const patchedListener = contextManager.bind(context, listener);
// store a weak reference of the user listener to ours
listeners.set(listener, patchedListener);

/**
* See comment at the start of this function for the explanation of this property.
*/
contextManager._wrapped = true;
try {
return original.call(this, event, patchedListener);
} finally {
contextManager._wrapped = false;
}
};
}

/**
*
*/
private _createPatchMap(ee: EventEmitter): PatchMap {
const map = Object.create(null);
// eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-member-access
(ee as any)[this._kOtListeners] = map;
return map;
}
/**
*
*/
private _getPatchMap(ee: EventEmitter): PatchMap | undefined {
return (ee as never)[this._kOtListeners];
}

/**
*
*/
// eslint-disable-next-line @typescript-eslint/no-explicit-any
private _bindFunction<T extends (...args: any[]) => any>(context: Context, target: T): T {
// eslint-disable-next-line @typescript-eslint/no-this-alias
const manager = this;
// eslint-disable-next-line @typescript-eslint/no-explicit-any
const contextWrapper = function (this: never, ...args: unknown[]): any {
return manager.with(context, () => target.apply(this, args));
};
Object.defineProperty(contextWrapper, 'length', {
enumerable: false,
configurable: true,
writable: false,
value: target.length,
});
/**
* It isn't possible to tell Typescript that contextWrapper is the same as T
* so we forced to cast as any here.
*/
// eslint-disable-next-line @typescript-eslint/no-explicit-any
return contextWrapper as any;
}

public abstract active(): Context;

public abstract with<A extends unknown[], F extends (...args: A) => ReturnType<F>>(
context: Context,
fn: F,
thisArg?: ThisParameterType<F>,
...args: A
): ReturnType<F>;

public abstract enable(): this;

public abstract disable(): this;
}

/**
*
*/
export class AsyncLocalStorageContextManager extends AbstractAsyncHooksContextManager {
private _asyncLocalStorage: AsyncLocalStorage<Context>;

public constructor() {
super();

// eslint-disable-next-line @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-explicit-any
const MaybeGlobalAsyncLocalStorage = (globalThis as any).AsyncLocalStorage;

if (!MaybeGlobalAsyncLocalStorage) {
DEBUG_BUILD &&
logger.warn(
"Tried to register AsyncLocalStorage async context strategy in a runtime that doesn't support AsyncLocalStorage.",
);
this._asyncLocalStorage = {
getStore() {
return undefined;
},
run(_store, callback, ...args) {
return callback.apply(this, args);
},
disable() {
// noop
},
};
} else {
this._asyncLocalStorage = new MaybeGlobalAsyncLocalStorage();
}
}

/**
*
*/
public active(): Context {
return this._asyncLocalStorage.getStore() ?? ROOT_CONTEXT;
}

/**
*
*/
public with<A extends unknown[], F extends (...args: A) => ReturnType<F>>(
context: Context,
fn: F,
thisArg?: ThisParameterType<F>,
...args: A
): ReturnType<F> {
const cb = thisArg == null ? fn : fn.bind(thisArg);
return this._asyncLocalStorage.run(context, cb as never, ...args);
}

/**
*
*/
public enable(): this {
return this;
}

/**
*
*/
public disable(): this {
this._asyncLocalStorage.disable();
return this;
}
}
Loading

0 comments on commit b7786ac

Please sign in to comment.