Skip to content
This repository was archived by the owner on Jan 21, 2026. It is now read-only.

Commit 132db9b

Browse files
kjinkjin
authored
fix: class-ify cls implementations (#708)
PR-URL: #708
1 parent 395a0c7 commit 132db9b

18 files changed

Lines changed: 862 additions & 375 deletions

‎src/cls-ah.ts‎

Lines changed: 0 additions & 145 deletions
This file was deleted.

‎src/cls.ts‎

Lines changed: 101 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/**
2-
* Copyright 2015 Google Inc. All Rights Reserved.
2+
* Copyright 2018 Google LLC
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -14,12 +14,17 @@
1414
* limitations under the License.
1515
*/
1616

17-
import * as CLS from 'continuation-local-storage';
17+
import {Logger} from '@google-cloud/common';
18+
import {EventEmitter} from 'events';
1819
import * as semver from 'semver';
1920

21+
import {AsyncHooksCLS} from './cls/async-hooks';
22+
import {AsyncListenerCLS} from './cls/async-listener';
23+
import {CLS, Func} from './cls/base';
24+
import {UniversalCLS} from './cls/universal';
2025
import {SpanDataType} from './constants';
21-
import {UNCORRELATED_SPAN, UNTRACED_SPAN} from './span-data';
2226
import {Trace, TraceSpan} from './trace';
27+
import {Singleton} from './util';
2328

2429
export interface RealRootContext {
2530
readonly span: TraceSpan;
@@ -43,78 +48,109 @@ export interface PhantomRootContext {
4348
*/
4449
export type RootContext = RealRootContext|PhantomRootContext;
4550

46-
export type Namespace = CLS.Namespace;
47-
export type Func<T> = CLS.Func<T>;
51+
const asyncHooksAvailable = semver.satisfies(process.version, '>=8');
4852

49-
const useAsyncHooks: boolean = semver.satisfies(process.version, '>=8') &&
50-
!!process.env.GCLOUD_TRACE_NEW_CONTEXT;
53+
export interface TraceCLSConfig { mechanism: 'async-listener'|'async-hooks'; }
5154

52-
const cls: typeof CLS =
53-
useAsyncHooks ? require('./cls-ah') : require('continuation-local-storage');
54-
55-
const TRACE_NAMESPACE = 'com.google.cloud.trace';
55+
export interface CLSConstructor {
56+
new(defaultContext: RootContext): CLS<RootContext>;
57+
}
5658

5759
/**
58-
* Stack traces are captured when a root span is started. Because the stack
59-
* trace height varies on the context propagation mechanism, to keep published
60-
* stack traces uniform we need to remove the top-most frames when using the
61-
* c-l-s module. Keep track of this number here.
60+
* An implementation of continuation-local storage for the Trace Agent.
61+
* In addition to the underlying API, there is a guarantee that when an instance
62+
* of this class is disabled, all context-manipulation methods will either be
63+
* no-ops or pass-throughs.
6264
*/
63-
export const ROOT_SPAN_STACK_OFFSET = useAsyncHooks ? 0 : 2;
65+
export class TraceCLS implements CLS<RootContext> {
66+
private currentCLS: CLS<RootContext>;
67+
// tslint:disable-next-line:variable-name CLSClass is a constructor.
68+
private CLSClass: CLSConstructor;
69+
private enabled = false;
70+
71+
private static UNCORRELATED: RootContext = {type: SpanDataType.UNCORRELATED};
72+
private static UNTRACED: RootContext = {type: SpanDataType.UNTRACED};
73+
74+
/**
75+
* Stack traces are captured when a root span is started. Because the stack
76+
* trace height varies on the context propagation mechanism, to keep published
77+
* stack traces uniform we need to remove the top-most frames when using the
78+
* c-l-s module. Keep track of this number here.
79+
*/
80+
readonly rootSpanStackOffset: number;
81+
82+
constructor(private readonly logger: Logger, config: TraceCLSConfig) {
83+
const useAH = config.mechanism === 'async-hooks' && asyncHooksAvailable;
84+
if (useAH) {
85+
this.CLSClass = AsyncHooksCLS;
86+
this.rootSpanStackOffset = 4;
87+
this.logger.info(
88+
'TraceCLS#constructor: Created [async-hooks] CLS instance.');
89+
} else {
90+
if (config.mechanism !== 'async-listener') {
91+
if (config.mechanism === 'async-hooks') {
92+
this.logger.error(
93+
'TraceCLS#constructor: [async-hooks]-based context',
94+
`propagation is not available in Node ${process.version}.`);
95+
} else {
96+
this.logger.error(
97+
'TraceCLS#constructor: The specified CLS mechanism',
98+
`[${config.mechanism}] was not recognized.`);
99+
}
100+
throw new Error(`CLS mechanism [${config.mechanism}] is invalid.`);
101+
}
102+
this.CLSClass = AsyncListenerCLS;
103+
this.rootSpanStackOffset = 8;
104+
this.logger.info(
105+
'TraceCLS#constructor: Created [async-listener] CLS instance.');
106+
}
107+
this.currentCLS = new UniversalCLS(TraceCLS.UNTRACED);
108+
this.currentCLS.enable();
109+
}
64110

65-
export function createNamespace(): CLS.Namespace {
66-
return cls.createNamespace(TRACE_NAMESPACE);
67-
}
111+
isEnabled(): boolean {
112+
return this.enabled;
113+
}
68114

69-
export function destroyNamespace(): void {
70-
cls.destroyNamespace(TRACE_NAMESPACE);
71-
}
115+
enable(): void {
116+
if (!this.enabled) {
117+
this.logger.info('TraceCLS#enable: Enabling CLS.');
118+
this.enabled = true;
119+
this.currentCLS.disable();
120+
this.currentCLS = new this.CLSClass(TraceCLS.UNCORRELATED);
121+
this.currentCLS.enable();
122+
}
123+
}
72124

73-
export function getNamespace(): CLS.Namespace {
74-
return cls.getNamespace(TRACE_NAMESPACE);
75-
}
125+
disable(): void {
126+
if (this.enabled) {
127+
this.logger.info('TraceCLS#disable: Disabling CLS.');
128+
this.enabled = false;
129+
this.currentCLS.disable();
130+
this.currentCLS = new UniversalCLS(TraceCLS.UNTRACED);
131+
this.currentCLS.enable();
132+
}
133+
}
76134

77-
/**
78-
* Get a RootContext object from continuation-local storage.
79-
*/
80-
export function getRootContext(): RootContext {
81-
// First getNamespace check is necessary in case any
82-
// patched closures escaped before the agent was stopped and the
83-
// namespace was destroyed.
84-
const namespace = getNamespace();
85-
if (namespace) {
86-
// A few things can be going on here:
87-
// 1. setRootContext has been called earlier to store a real root span
88-
// in continuation-local storage, so retrieve it.
89-
// 2. setRootContext has been called earlier to explicitly specify that
90-
// the request corresponding to this continuation is _not_ being traced
91-
// (by being passed UNTRACED_SPAN), so retrieve it as well.
92-
// 3. setRootContext has _never_ been called in this continuation. This
93-
// indicates that context was lost, and namespace.get('root') will
94-
// return null. Therefore, explicitly return UNCORRELATED_SPAN to
95-
// indicate that context was lost.
96-
return namespace.get('root') || UNCORRELATED_SPAN;
97-
} else {
98-
// No namespace indicates that the Trace Agent is disabled. This is a
99-
// special case where _all_ requests are explicitly not being traced,
100-
// so return UNTRACED_SPAN to be consistent with that.
101-
return UNTRACED_SPAN;
135+
getContext(): RootContext {
136+
return this.currentCLS.getContext();
102137
}
103-
}
104138

105-
/**
106-
* Store a RootContext object in continuation-local storage.
107-
* @param rootContext Either a root span or UNTRACED_SPAN. It doesn't make
108-
* sense to pass UNCORRELATED_SPAN, which is a value specifically reserved for
109-
* when getRootContext is known to give an unusable value.
110-
*/
111-
export function setRootContext(rootContext: RootContext): void {
112-
getNamespace().set('root', rootContext);
113-
}
139+
setContext(value: RootContext): void {
140+
this.currentCLS.setContext(value);
141+
}
114142

115-
// This is only used in tests (and is temporary), so it doesn't apply in the
116-
// comment in getRootContext about the possible values of namespace.get('root').
117-
// It's functionally identical to setRootContext(null).
118-
export function clearRootContext(): void {
119-
setRootContext(UNCORRELATED_SPAN);
143+
runWithNewContext<T>(fn: Func<T>): T {
144+
return this.currentCLS.runWithNewContext(fn);
145+
}
146+
147+
bindWithCurrentContext<T>(fn: Func<T>): Func<T> {
148+
return this.currentCLS.bindWithCurrentContext(fn);
149+
}
150+
151+
patchEmitterToPropagateContext<T>(ee: EventEmitter): void {
152+
this.currentCLS.patchEmitterToPropagateContext(ee);
153+
}
120154
}
155+
156+
export const cls = new Singleton(TraceCLS);

0 commit comments

Comments
 (0)