Skip to content

Commit 93b8fcb

Browse files
alkatrivedialkatrivedi
authored
Merge branch 'main' into uuid-support
2 parents 0824f39 + 0008038 commit 93b8fcb

File tree

4 files changed

+505
-0
lines changed

4 files changed

+505
-0
lines changed

‎.gitignore

+1
Original file line numberDiff line numberDiff line change
@@ -12,3 +12,4 @@ system-test/*key.json
1212
.DS_Store
1313
package-lock.json
1414
__pycache__
15+
.idea

‎src/request_id_header.ts

+200
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,200 @@
1+
/**
2+
* Copyright 2025 Google LLC. All Rights Reserved.
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 {randomBytes} from 'crypto';
18+
// eslint-disable-next-line n/no-extraneous-import
19+
import * as grpc from '@grpc/grpc-js';
20+
const randIdForProcess = randomBytes(8)
21+
.readUint32LE(0)
22+
.toString(16)
23+
.padStart(8, '0');
24+
const X_GOOG_SPANNER_REQUEST_ID_HEADER = 'x-goog-spanner-request-id';
25+
26+
class AtomicCounter {
27+
private readonly backingBuffer: Uint32Array;
28+
29+
constructor(initialValue?: number) {
30+
this.backingBuffer = new Uint32Array(
31+
new SharedArrayBuffer(Uint32Array.BYTES_PER_ELEMENT)
32+
);
33+
if (initialValue) {
34+
this.increment(initialValue);
35+
}
36+
}
37+
38+
public increment(n?: number): number {
39+
if (!n) {
40+
n = 1;
41+
}
42+
Atomics.add(this.backingBuffer, 0, n);
43+
return this.value();
44+
}
45+
46+
public value(): number {
47+
return Atomics.load(this.backingBuffer, 0);
48+
}
49+
50+
public toString(): string {
51+
return `${this.value()}`;
52+
}
53+
54+
public reset() {
55+
Atomics.store(this.backingBuffer, 0, 0);
56+
}
57+
}
58+
59+
const REQUEST_HEADER_VERSION = 1;
60+
61+
function craftRequestId(
62+
nthClientId: number,
63+
channelId: number,
64+
nthRequest: number,
65+
attempt: number
66+
) {
67+
return `${REQUEST_HEADER_VERSION}.${randIdForProcess}.${nthClientId}.${channelId}.${nthRequest}.${attempt}`;
68+
}
69+
70+
const nthClientId = new AtomicCounter();
71+
72+
// Only exported for deterministic testing.
73+
export function resetNthClientId() {
74+
nthClientId.reset();
75+
}
76+
77+
/*
78+
* nextSpannerClientId increments the internal
79+
* counter for created SpannerClients, for use
80+
* with x-goog-spanner-request-id.
81+
*/
82+
function nextSpannerClientId(): number {
83+
nthClientId.increment(1);
84+
return nthClientId.value();
85+
}
86+
87+
function newAtomicCounter(n?: number): AtomicCounter {
88+
return new AtomicCounter(n);
89+
}
90+
91+
interface withHeaders {
92+
headers: {[k: string]: string};
93+
}
94+
95+
function extractRequestID(config: any): string {
96+
if (!config) {
97+
return '';
98+
}
99+
100+
const hdrs = config as withHeaders;
101+
if (hdrs && hdrs.headers) {
102+
return hdrs.headers[X_GOOG_SPANNER_REQUEST_ID_HEADER];
103+
}
104+
return '';
105+
}
106+
107+
function injectRequestIDIntoError(config: any, err: Error) {
108+
if (!err) {
109+
return;
110+
}
111+
112+
// Inject that RequestID into the actual
113+
// error object regardless of the type.
114+
const requestID = extractRequestID(config);
115+
if (requestID) {
116+
Object.assign(err, {requestID: requestID});
117+
}
118+
}
119+
120+
interface withNextNthRequest {
121+
_nextNthRequest: Function;
122+
}
123+
124+
interface withMetadataWithRequestId {
125+
_nthClientId: number;
126+
_channelId: number;
127+
}
128+
129+
function injectRequestIDIntoHeaders(
130+
headers: {[k: string]: string},
131+
session: any,
132+
nthRequest?: number,
133+
attempt?: number
134+
) {
135+
if (!session) {
136+
return headers;
137+
}
138+
139+
if (!nthRequest) {
140+
const database = session.parent as withNextNthRequest;
141+
if (!(database && typeof database._nextNthRequest === 'function')) {
142+
return headers;
143+
}
144+
nthRequest = database._nextNthRequest();
145+
}
146+
147+
attempt = attempt || 1;
148+
return _metadataWithRequestId(session, nthRequest!, attempt, headers);
149+
}
150+
151+
function _metadataWithRequestId(
152+
session: any,
153+
nthRequest: number,
154+
attempt: number,
155+
priorMetadata?: {[k: string]: string}
156+
): {[k: string]: string} {
157+
if (!priorMetadata) {
158+
priorMetadata = {};
159+
}
160+
const withReqId = {
161+
...priorMetadata,
162+
};
163+
const database = session.parent as withMetadataWithRequestId;
164+
let clientId = 1;
165+
let channelId = 1;
166+
if (database) {
167+
clientId = database._nthClientId || 1;
168+
channelId = database._channelId || 1;
169+
}
170+
withReqId[X_GOOG_SPANNER_REQUEST_ID_HEADER] = craftRequestId(
171+
clientId,
172+
channelId,
173+
nthRequest,
174+
attempt
175+
);
176+
return withReqId;
177+
}
178+
179+
function nextNthRequest(database): number {
180+
if (!(database && typeof database._nextNthRequest === 'function')) {
181+
return 1;
182+
}
183+
return database._nextNthRequest();
184+
}
185+
186+
export interface RequestIDError extends grpc.ServiceError {
187+
requestID: string;
188+
}
189+
190+
export {
191+
AtomicCounter,
192+
X_GOOG_SPANNER_REQUEST_ID_HEADER,
193+
craftRequestId,
194+
injectRequestIDIntoError,
195+
injectRequestIDIntoHeaders,
196+
nextNthRequest,
197+
nextSpannerClientId,
198+
newAtomicCounter,
199+
randIdForProcess,
200+
};

‎test/request_id_header.ts

+171
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,171 @@
1+
/**
2+
* Copyright 2025 Google LLC. All Rights Reserved.
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+
/* eslint-disable prefer-rest-params */
18+
import * as assert from 'assert';
19+
import {
20+
RequestIDError,
21+
X_GOOG_SPANNER_REQUEST_ID_HEADER,
22+
craftRequestId,
23+
injectRequestIDIntoError,
24+
injectRequestIDIntoHeaders,
25+
newAtomicCounter,
26+
nextNthRequest,
27+
randIdForProcess,
28+
} from '../src/request_id_header';
29+
30+
describe('RequestId', () => {
31+
describe('AtomicCounter', () => {
32+
it('Constructor with initialValue', done => {
33+
const ac0 = newAtomicCounter();
34+
assert.deepStrictEqual(ac0.value(), 0);
35+
assert.deepStrictEqual(
36+
ac0.increment(2),
37+
2,
38+
'increment should return the added value'
39+
);
40+
assert.deepStrictEqual(
41+
ac0.value(),
42+
2,
43+
'increment should have modified the value'
44+
);
45+
46+
const ac1 = newAtomicCounter(1);
47+
assert.deepStrictEqual(ac1.value(), 1);
48+
assert.deepStrictEqual(
49+
ac1.increment(1 << 27),
50+
(1 << 27) + 1,
51+
'increment should return the added value'
52+
);
53+
assert.deepStrictEqual(
54+
ac1.value(),
55+
(1 << 27) + 1,
56+
'increment should have modified the value'
57+
);
58+
done();
59+
});
60+
61+
it('reset', done => {
62+
const ac0 = newAtomicCounter(1);
63+
ac0.increment();
64+
assert.strictEqual(ac0.value(), 2);
65+
ac0.reset();
66+
assert.strictEqual(ac0.value(), 0);
67+
done();
68+
});
69+
70+
it('toString', done => {
71+
const ac0 = newAtomicCounter(1);
72+
ac0.increment();
73+
assert.strictEqual(ac0.value(), 2);
74+
assert.strictEqual(ac0.toString(), '2');
75+
assert.strictEqual(`${ac0}`, '2');
76+
done();
77+
});
78+
});
79+
80+
describe('craftRequestId', () => {
81+
it('has a 32-bit hex-formatted process-id', done => {
82+
assert.match(
83+
randIdForProcess,
84+
/^[0-9A-Fa-f]{8}$/,
85+
`process-id should be a 32-bit hexadecimal number, but was ${randIdForProcess}`
86+
);
87+
done();
88+
});
89+
90+
it('with attempts', done => {
91+
assert.strictEqual(
92+
craftRequestId(1, 2, 3, 4),
93+
`1.${randIdForProcess}.1.2.3.4`
94+
);
95+
done();
96+
});
97+
});
98+
99+
describe('injectRequestIDIntoError', () => {
100+
it('with non-null error', done => {
101+
const err: Error = new Error('this one');
102+
const config = {headers: {}};
103+
config.headers[X_GOOG_SPANNER_REQUEST_ID_HEADER] = '1.2.3.4.5.6';
104+
injectRequestIDIntoError(config, err);
105+
assert.strictEqual((err as RequestIDError).requestID, '1.2.3.4.5.6');
106+
done();
107+
});
108+
});
109+
110+
describe('injectRequestIDIntoHeaders', () => {
111+
it('with null session', done => {
112+
const hdrs = {};
113+
injectRequestIDIntoHeaders(hdrs, null, 2, 1);
114+
done();
115+
});
116+
117+
it('with nthRequest explicitly passed in', done => {
118+
const session = {
119+
parent: {
120+
_nextNthRequest: () => {
121+
return 5;
122+
},
123+
},
124+
};
125+
const got = injectRequestIDIntoHeaders({}, session, 2, 5);
126+
const want = {
127+
'x-goog-spanner-request-id': `1.${randIdForProcess}.1.1.2.5`,
128+
};
129+
assert.deepStrictEqual(got, want);
130+
done();
131+
});
132+
133+
it('infer nthRequest from session', done => {
134+
const session = {
135+
parent: {
136+
_nextNthRequest: () => {
137+
return 5;
138+
},
139+
},
140+
};
141+
142+
const inputHeaders: {[k: string]: string} = {};
143+
const got = injectRequestIDIntoHeaders(inputHeaders, session);
144+
const want = {
145+
'x-goog-spanner-request-id': `1.${randIdForProcess}.1.1.5.1`,
146+
};
147+
assert.deepStrictEqual(got, want);
148+
done();
149+
});
150+
});
151+
152+
describe('nextNthRequest', () => {
153+
const fauxDatabase = {};
154+
assert.deepStrictEqual(
155+
nextNthRequest(fauxDatabase),
156+
1,
157+
'Without override, should default to 1'
158+
);
159+
160+
Object.assign(fauxDatabase, {
161+
_nextNthRequest: () => {
162+
return 4;
163+
},
164+
});
165+
assert.deepStrictEqual(
166+
nextNthRequest(fauxDatabase),
167+
4,
168+
'With override should infer value'
169+
);
170+
});
171+
});

0 commit comments

Comments
 (0)