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

Commit 5d000e9

Browse files
kjinkjin
authored
feat: allow "disabling" cls, and relax requirements for creating root spans (#728)
PR-URL: #728
1 parent edb8135 commit 5d000e9

15 files changed

Lines changed: 241 additions & 85 deletions

‎.circleci/config.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,7 @@ postgres_service: &postgres_service
7474
POSTGRES_DB: test
7575

7676
mysql_service: &mysql_service
77-
image: mysql
77+
image: mysql:5
7878
environment:
7979
MYSQL_ROOT_PASSWORD: Password12!
8080
MYSQL_DATABASE: test

‎src/cls.ts‎

Lines changed: 49 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@ import {SpanDataType} from './constants';
2626
import {Trace, TraceSpan} from './trace';
2727
import {Singleton} from './util';
2828

29+
const asyncHooksAvailable = semver.satisfies(process.version, '>=8');
30+
2931
export interface RealRootContext {
3032
readonly span: TraceSpan;
3133
readonly trace: Trace;
@@ -48,11 +50,34 @@ export interface PhantomRootContext {
4850
*/
4951
export type RootContext = RealRootContext|PhantomRootContext;
5052

51-
const asyncHooksAvailable = semver.satisfies(process.version, '>=8');
53+
/**
54+
* An enumeration of the possible mechanisms for supporting context propagation
55+
* through continuation-local storage.
56+
*/
57+
export enum TraceCLSMechanism {
58+
/**
59+
* Use the AsyncHooksCLS class to propagate root span context.
60+
* Only available in Node 8+.
61+
*/
62+
ASYNC_HOOKS = 'async-hooks',
63+
/**
64+
* Use the AsyncListenerCLS class to propagate root span context.
65+
* Note that continuation-local-storage should be loaded as the first module.
66+
*/
67+
ASYNC_LISTENER = 'async-listener',
68+
/**
69+
* Do not use any special mechanism to propagate root span context.
70+
* Only a single root span can be open at a time.
71+
*/
72+
NONE = 'none'
73+
}
5274

53-
export interface TraceCLSConfig { mechanism: 'async-listener'|'async-hooks'; }
75+
/**
76+
* Configuration options passed to the TraceCLS constructor.
77+
*/
78+
export interface TraceCLSConfig { mechanism: TraceCLSMechanism; }
5479

55-
export interface CLSConstructor {
80+
interface CLSConstructor {
5681
new(defaultContext: RootContext): CLS<RootContext>;
5782
}
5883

@@ -80,30 +105,29 @@ export class TraceCLS implements CLS<RootContext> {
80105
readonly rootSpanStackOffset: number;
81106

82107
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.`);
108+
switch (config.mechanism) {
109+
case TraceCLSMechanism.ASYNC_HOOKS:
110+
if (!asyncHooksAvailable) {
111+
throw new Error(`CLS mechanism [${
112+
config.mechanism}] is not compatible with Node <8.`);
99113
}
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.');
114+
this.CLSClass = AsyncHooksCLS;
115+
this.rootSpanStackOffset = 4;
116+
break;
117+
case TraceCLSMechanism.ASYNC_LISTENER:
118+
this.CLSClass = AsyncListenerCLS;
119+
this.rootSpanStackOffset = 8;
120+
break;
121+
case TraceCLSMechanism.NONE:
122+
this.CLSClass = UniversalCLS;
123+
this.rootSpanStackOffset = 4;
124+
break;
125+
default:
126+
throw new Error(
127+
`CLS mechanism [${config.mechanism}] was not recognized.`);
106128
}
129+
this.logger.info(
130+
`TraceCLS#constructor: Created [${config.mechanism}] CLS instance.`);
107131
this.currentCLS = new UniversalCLS(TraceCLS.UNTRACED);
108132
this.currentCLS.enable();
109133
}

‎src/cls/async-listener.ts‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,6 @@ export class AsyncListenerCLS<Context extends {}> implements CLS<Context> {
7171

7272
runWithNewContext<T>(fn: Func<T>): T {
7373
return this.getNamespace().runAndReturn(() => {
74-
this.setContext(this.defaultContext);
7574
return fn();
7675
});
7776
}

‎src/cls/universal.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ export class UniversalCLS<Context> implements CLS<Context> {
4040

4141
disable(): void {
4242
this.enabled = false;
43-
this.currentContext = this.defaultContext;
43+
this.setContext(this.defaultContext);
4444
}
4545

4646
getContext(): Context {

‎src/config.ts‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,22 @@ import * as path from 'path';
1919
const pluginDirectory =
2020
path.join(path.resolve(__dirname, '..'), 'src', 'plugins');
2121

22+
export type CLSMechanism = 'none'|'auto';
23+
2224
/** Available configuration options. */
2325
export interface Config {
26+
/**
27+
* The trace context propagation mechanism to use. The following options are
28+
* available:
29+
* - 'auto' uses continuation-local-storage, unless async_hooks is available
30+
* _and_ the environment variable GCLOUD_TRACE_NEW_CONTEXT is set, in which
31+
* case async_hooks will be used instead.
32+
* - 'none' disables CLS completely.
33+
* The 'auto' mechanism is used by default if this configuration option is
34+
* not explicitly set.
35+
*/
36+
clsMechanism?: CLSMechanism;
37+
2438
/**
2539
* Log levels: 0=disabled, 1=error, 2=warn, 3=info, 4=debug
2640
* The value of GCLOUD_TRACE_LOGLEVEL takes precedence over this value.
@@ -166,6 +180,7 @@ export interface Config {
166180
* user-provided value will be used to extend the default value.
167181
*/
168182
export const defaultConfig = {
183+
clsMechanism: 'auto' as CLSMechanism,
169184
logLevel: 1,
170185
enabled: true,
171186
enhancedDatabaseReporting: false,

‎src/index.ts‎

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -27,9 +27,9 @@ if (!useAH) {
2727
}
2828

2929
import * as common from '@google-cloud/common';
30-
import {cls} from './cls';
30+
import {cls, TraceCLSConfig, TraceCLSMechanism} from './cls';
3131
import {Constants} from './constants';
32-
import {Config, defaultConfig} from './config';
32+
import {Config, defaultConfig, CLSMechanism} from './config';
3333
import * as extend from 'extend';
3434
import * as path from 'path';
3535
import * as PluginTypes from './plugin-types';
@@ -56,6 +56,7 @@ for (let i = 0; i < filesLoadedBeforeTrace.length; i++) {
5656
interface TopLevelConfig {
5757
enabled: boolean;
5858
logLevel: number;
59+
clsMechanism: CLSMechanism;
5960
}
6061

6162
// PluginLoaderConfig extends TraceAgentConfig
@@ -178,12 +179,14 @@ export function start(projectConfig?: Config): PluginTypes.TraceAgent {
178179

179180
try {
180181
// Initialize context propagation mechanism.
181-
// TODO(kjin): Publicly expose this field.
182-
cls.create(logger, {
183-
mechanism: useAH ? 'async-hooks' : 'async-listener',
184-
[FORCE_NEW]: config[FORCE_NEW]
185-
})
186-
.enable();
182+
const m = config.clsMechanism;
183+
const clsConfig: Forceable<TraceCLSConfig> = {
184+
mechanism: m === 'auto' ? (useAH ? TraceCLSMechanism.ASYNC_HOOKS :
185+
TraceCLSMechanism.ASYNC_LISTENER) :
186+
m as TraceCLSMechanism,
187+
[FORCE_NEW]: config[FORCE_NEW]
188+
};
189+
cls.create(logger, clsConfig).enable();
187190

188191
traceWriter.create(logger, config).initialize((err) => {
189192
if (err) {

‎src/trace-api.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -128,7 +128,8 @@ export class TraceAgent implements TraceAgentInterface {
128128

129129
// TODO validate options
130130
// Don't create a root span if we are already in a root span
131-
if (cls.get().getContext().type === SpanDataType.ROOT) {
131+
const rootSpan = cls.get().getContext();
132+
if (rootSpan.type === SpanDataType.ROOT && !rootSpan.span.endTime) {
132133
this.logger!.warn(`TraceApi#runInRootSpan: [${
133134
this.pluginName}] Cannot create nested root spans.`);
134135
return fn(UNCORRELATED_SPAN);

‎test/plugins/test-trace-google-gax.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,11 +59,17 @@ describe('Tracing with google-gax', () => {
5959
let googleGax: GaxModule;
6060

6161
before(() => {
62+
trace.setCLS();
6263
trace.setPluginLoader();
6364
trace.start();
6465
googleGax = require('./fixtures/google-gax0.16');
6566
});
6667

68+
after(() => {
69+
trace.setCLS(trace.TestCLS);
70+
trace.setPluginLoader(trace.TestPluginLoader);
71+
});
72+
6773
it(`doesn't break context`, (done) => {
6874
const authPromise = Promise.resolve(
6975
((args, metadata, opts, cb) => {

‎test/test-cls.ts‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ import {ITestDefinition} from 'mocha';
2424
import * as semver from 'semver';
2525
import {inspect} from 'util';
2626

27-
import {TraceCLS, TraceCLSConfig} from '../src/cls';
27+
import {TraceCLS, TraceCLSConfig, TraceCLSMechanism} from '../src/cls';
2828
import {AsyncHooksCLS} from '../src/cls/async-hooks';
2929
import {AsyncListenerCLS} from '../src/cls/async-listener';
3030
import {CLS} from '../src/cls/base';
@@ -221,8 +221,15 @@ describe('Continuation-Local Storage', () => {
221221

222222
describe('TraceCLS', () => {
223223
const validTestCases: TraceCLSConfig[] = asyncAwaitSupported ?
224-
[{mechanism: 'async-hooks'}, {mechanism: 'async-listener'}] :
225-
[{mechanism: 'async-listener'}];
224+
[
225+
{mechanism: TraceCLSMechanism.ASYNC_HOOKS},
226+
{mechanism: TraceCLSMechanism.ASYNC_LISTENER},
227+
{mechanism: TraceCLSMechanism.NONE}
228+
] :
229+
[
230+
{mechanism: TraceCLSMechanism.ASYNC_LISTENER},
231+
{mechanism: TraceCLSMechanism.NONE}
232+
];
226233
for (const testCase of validTestCases) {
227234
describe(`with configuration ${inspect(testCase)}`, () => {
228235
const logger = new TestLogger();

‎test/test-config-cls.ts‎

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
/**
2+
* Copyright 2018 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
import {Logger} from '@google-cloud/common';
18+
import * as assert from 'assert';
19+
import * as semver from 'semver';
20+
import * as util from 'util';
21+
22+
import {TraceCLSConfig, TraceCLSMechanism} from '../src/cls';
23+
24+
import * as trace from './trace';
25+
26+
describe('Behavior set by config for context propagation mechanism', () => {
27+
const useAH = semver.satisfies(process.version, '>=8') &&
28+
!!process.env.GCLOUD_TRACE_NEW_CONTEXT;
29+
const autoMechanism =
30+
useAH ? TraceCLSMechanism.ASYNC_HOOKS : TraceCLSMechanism.ASYNC_LISTENER;
31+
let capturedConfig: TraceCLSConfig|null;
32+
33+
class CaptureConfigTestCLS extends trace.TestCLS {
34+
constructor(logger: Logger, config: TraceCLSConfig) {
35+
super(logger, config);
36+
// Capture the config object passed into this constructor.
37+
capturedConfig = config;
38+
}
39+
}
40+
41+
beforeEach(() => {
42+
capturedConfig = null;
43+
});
44+
45+
before(() => {
46+
trace.setCLS(CaptureConfigTestCLS);
47+
});
48+
49+
after(() => {
50+
trace.setCLS(trace.TestCLS);
51+
});
52+
53+
const testCases: Array<
54+
{tracingConfig: trace.Config, contextPropagationConfig: TraceCLSConfig}> =
55+
[
56+
{
57+
tracingConfig: {clsMechanism: 'none'},
58+
contextPropagationConfig: {mechanism: 'none'}
59+
},
60+
{
61+
tracingConfig: {clsMechanism: 'auto'},
62+
contextPropagationConfig: {mechanism: autoMechanism}
63+
},
64+
{
65+
tracingConfig: {},
66+
contextPropagationConfig: {mechanism: autoMechanism}
67+
},
68+
{
69+
// tslint:disable:no-any
70+
tracingConfig: {clsMechanism: 'unknown' as any},
71+
contextPropagationConfig: {mechanism: 'unknown' as any}
72+
// tslint:enable:no-any
73+
}
74+
];
75+
76+
for (const testCase of testCases) {
77+
it(`should be as expected for config: ${
78+
util.inspect(testCase.tracingConfig)}`,
79+
() => {
80+
trace.start(testCase.tracingConfig);
81+
assert.ok(capturedConfig);
82+
assert.strictEqual(
83+
capturedConfig!.mechanism,
84+
testCase.contextPropagationConfig.mechanism);
85+
});
86+
}
87+
});

0 commit comments

Comments
 (0)