/** * Copyright 2018 Google LLC * * 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 * * http://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. */ // This module requires continuation-local-storage in the AsyncListenerCLS // constructor, rather than upon module load. import * as clsModule from 'continuation-local-storage'; import { EventEmitter } from 'events'; import { CLS, Func } from './base'; type CLSModule = typeof clsModule; /** * An implementation of continuation-local storage that wraps the * "continuation-local-storage" module. */ export class AsyncListenerCLS implements CLS { static readonly TRACE_NAMESPACE = 'com.google.cloud.trace'; static readonly ROOT_CONTEXT_KEY = 'root'; private readonly cls: CLSModule; private readonly defaultContext: Context; constructor(defaultContext: Context) { // Conditionally load continuation-local-storage. // We make this a member field instead of assigning it to a module-scope // object to make its access uncomplicated. this.cls = require('continuation-local-storage'); this.defaultContext = defaultContext; } isEnabled(): boolean { return !!this.getNamespace(); } enable(): void { this.cls.createNamespace(AsyncListenerCLS.TRACE_NAMESPACE); } disable(): void { this.cls.destroyNamespace(AsyncListenerCLS.TRACE_NAMESPACE); } private getNamespace(): clsModule.Namespace { return this.cls.getNamespace(AsyncListenerCLS.TRACE_NAMESPACE); } getContext(): Context { const result = this.getNamespace().get(AsyncListenerCLS.ROOT_CONTEXT_KEY); if (result) { return result; } return this.defaultContext; } runWithContext(fn: Func, value: Context): T { const namespace = this.getNamespace(); return namespace.runAndReturn(() => { namespace.set(AsyncListenerCLS.ROOT_CONTEXT_KEY, value); return fn(); }); } bindWithCurrentContext(fn: Func): Func { return this.getNamespace().bind(fn) as Func; } patchEmitterToPropagateContext(ee: EventEmitter): void { return this.getNamespace().bindEmitter(ee); } }