feat(x-goog-spanner-request-id): add bases

src/request_id_header.ts

@@ -0,0 +1,197 @@
+/**
+ * Copyright 2025 Google LLC. All Rights Reserved.
+ *
+ * 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.
+ */
+
+import {randomBytes} from 'crypto';
+// eslint-disable-next-line n/no-extraneous-import
+import * as grpc from '@grpc/grpc-js';
+const randIdForProcess = randomBytes(8).readBigUint64LE(0).toString();
+const X_GOOG_SPANNER_REQUEST_ID_HEADER = 'x-goog-spanner-request-id';
+
+class AtomicCounter {
+  private readonly backingBuffer: Uint32Array;
+
+  constructor(initialValue?: number) {
+    this.backingBuffer = new Uint32Array(
+      new SharedArrayBuffer(Uint32Array.BYTES_PER_ELEMENT)
+    );
+    if (initialValue) {
+      this.increment(initialValue);
+    }
+  }
+
+  public increment(n?: number): number {
+    if (!n) {
+      n = 1;
+    }
+    Atomics.add(this.backingBuffer, 0, n);
+    return this.value();
+  }
+
+  public value(): number {
+    return Atomics.load(this.backingBuffer, 0);
+  }
+
+  public toString(): string {
+    return `${this.value()}`;
+  }
+
+  public reset() {
+    Atomics.store(this.backingBuffer, 0, 0);
+  }
+}
+
+const REQUEST_HEADER_VERSION = 1;
+
+function craftRequestId(
+  nthClientId: number,
+  channelId: number,
+  nthRequest: number,
+  attempt: number
+) {
+  return `${REQUEST_HEADER_VERSION}.${randIdForProcess}.${nthClientId}.${channelId}.${nthRequest}.${attempt}`;
+}
+
+const nthClientId = new AtomicCounter();
+
+// Only exported for deterministic testing.
+export function resetNthClientId() {
+  nthClientId.reset();
+}
+
+/*
+ * nextSpannerClientId increments the internal
+ * counter for created SpannerClients, for use
+ * with x-goog-spanner-request-id.
+ */
+function nextSpannerClientId(): number {
+  nthClientId.increment(1);
+  return nthClientId.value();
+}
+
+function newAtomicCounter(n?: number): AtomicCounter {
+  return new AtomicCounter(n);
+}
+
+interface withHeaders {
+  headers: {[k: string]: string};
+}
+
+function extractRequestID(config: any): string {
+  if (!config) {
+    return '';
+  }
+
+  const hdrs = config as withHeaders;
+  if (hdrs && hdrs.headers) {
+    return hdrs.headers[X_GOOG_SPANNER_REQUEST_ID_HEADER];
+  }
+  return '';
+}
+
+function injectRequestIDIntoError(config: any, err: Error) {
+  if (!err) {
+    return;
+  }
+
+  // Inject that RequestID into the actual
+  // error object regardless of the type.
+  const requestID = extractRequestID(config);
+  if (requestID) {
+    Object.assign(err, {requestID: requestID});
+  }
+}
+
+interface withNextNthRequest {
+  _nextNthRequest: Function;
+}
+
+interface withMetadataWithRequestId {
+  _nthClientId: number;
+  _channelId: number;
+}
+
+function injectRequestIDIntoHeaders(
+  headers: {[k: string]: string},
+  session: any,
+  nthRequest?: number,
+  attempt?: number
+) {
+  if (!session) {
+    return headers;
+  }
+
+  if (!nthRequest) {
+    const database = session.parent as withNextNthRequest;
+    if (!(database && typeof database._nextNthRequest === 'function')) {
+      return headers;
+    }
+    nthRequest = database._nextNthRequest();
+  }
+
+  attempt = attempt || 1;
+  return _metadataWithRequestId(session, nthRequest!, attempt, headers);
+}
+
+function _metadataWithRequestId(
+  session: any,
+  nthRequest: number,
+  attempt: number,
+  priorMetadata?: {[k: string]: string}
+): {[k: string]: string} {
+  if (!priorMetadata) {
+    priorMetadata = {};
+  }
+  const withReqId = {
+    ...priorMetadata,
+  };
+  const database = session.parent as withMetadataWithRequestId;
+  let clientId = 1;
+  let channelId = 1;
+  if (database) {
+    clientId = database._nthClientId || 1;
+    channelId = database._channelId || 1;
+  }
+  withReqId[X_GOOG_SPANNER_REQUEST_ID_HEADER] = craftRequestId(
+    clientId,
+    channelId,
+    nthRequest,
+    attempt
+  );
+  return withReqId;
+}
+
+function nextNthRequest(database): number {
+  if (!(database && typeof database._nextNthRequest === 'function')) {
+    return 1;
+  }
+  return database._nextNthRequest();
+}
+
+export interface RequestIDError extends grpc.ServiceError {
+  requestID: string;
+}
+
+export {
+  AtomicCounter,
+  X_GOOG_SPANNER_REQUEST_ID_HEADER,
+  craftRequestId,
+  injectRequestIDIntoError,
+  injectRequestIDIntoHeaders,
+  nextNthRequest,
+  nextSpannerClientId,
+  newAtomicCounter,
+  randIdForProcess,
+};

test/request_id_header.ts

@@ -0,0 +1,162 @@
+/**
+ * Copyright 2025 Google LLC. All Rights Reserved.
+ *
+ * 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.
+ */
+
+/* eslint-disable prefer-rest-params */
+import * as assert from 'assert';
+import {
+  RequestIDError,
+  X_GOOG_SPANNER_REQUEST_ID_HEADER,
+  craftRequestId,
+  injectRequestIDIntoError,
+  injectRequestIDIntoHeaders,
+  newAtomicCounter,
+  nextNthRequest,
+  randIdForProcess,
+} from '../src/request_id_header';
+
+describe('RequestId', () => {
+  describe('AtomicCounter', () => {
+    it('Constructor with initialValue', done => {
+      const ac0 = newAtomicCounter();
+      assert.deepStrictEqual(ac0.value(), 0);
+      assert.deepStrictEqual(
+        ac0.increment(2),
+        2,
+        'increment should return the added value'
+      );
+      assert.deepStrictEqual(
+        ac0.value(),
+        2,
+        'increment should have modified the value'
+      );
+
+      const ac1 = newAtomicCounter(1);
+      assert.deepStrictEqual(ac1.value(), 1);
+      assert.deepStrictEqual(
+        ac1.increment(1 << 27),
+        (1 << 27) + 1,
+        'increment should return the added value'
+      );
+      assert.deepStrictEqual(
+        ac1.value(),
+        (1 << 27) + 1,
+        'increment should have modified the value'
+      );
+      done();
+    });
+
+    it('reset', done => {
+      const ac0 = newAtomicCounter(1);
+      ac0.increment();
+      assert.strictEqual(ac0.value(), 2);
+      ac0.reset();
+      assert.strictEqual(ac0.value(), 0);
+      done();
+    });
+
+    it('toString', done => {
+      const ac0 = newAtomicCounter(1);
+      ac0.increment();
+      assert.strictEqual(ac0.value(), 2);
+      assert.strictEqual(ac0.toString(), '2');
+      assert.strictEqual(`${ac0}`, '2');
+      done();
+    });
+  });
+
+  describe('craftRequestId', () => {
+    it('with attempts', done => {
+      assert.strictEqual(
+        craftRequestId(1, 2, 3, 4),
+        `1.${randIdForProcess}.1.2.3.4`
+      );
+      done();
+    });
+  });
+
+  describe('injectRequestIDIntoError', () => {
+    it('with non-null error', done => {
+      const err: Error = new Error('this one');
+      const config = {headers: {}};
+      config.headers[X_GOOG_SPANNER_REQUEST_ID_HEADER] = '1.2.3.4.5.6';
+      injectRequestIDIntoError(config, err);
+      assert.strictEqual((err as RequestIDError).requestID, '1.2.3.4.5.6');
+      done();
+    });
+  });
+
+  describe('injectRequestIDIntoHeaders', () => {
+    it('with null session', done => {
+      const hdrs = {};
+      injectRequestIDIntoHeaders(hdrs, null, 2, 1);
+      done();
+    });
+
+    it('with nthRequest explicitly passed in', done => {
+      const session = {
+        parent: {
+          _nextNthRequest: () => {
+            return 5;
+          },
+        },
+      };
+      const got = injectRequestIDIntoHeaders({}, session, 2, 5);
+      const want = {
+        'x-goog-spanner-request-id': `1.${randIdForProcess}.1.1.2.5`,
+      };
+      assert.deepStrictEqual(got, want);
+      done();
+    });
+
+    it('infer nthRequest from session', done => {
+      const session = {
+        parent: {
+          _nextNthRequest: () => {
+            return 5;
+          },
+        },
+      };
+
+      const inputHeaders: {[k: string]: string} = {};
+      const got = injectRequestIDIntoHeaders(inputHeaders, session);
+      const want = {
+        'x-goog-spanner-request-id': `1.${randIdForProcess}.1.1.5.1`,
+      };
+      assert.deepStrictEqual(got, want);
+      done();
+    });
+  });
+
+  describe('nextNthRequest', () => {
+    const fauxDatabase = {};
+    assert.deepStrictEqual(
+      nextNthRequest(fauxDatabase),
+      1,
+      'Without override, should default to 1'
+    );
+
+    Object.assign(fauxDatabase, {
+      _nextNthRequest: () => {
+        return 4;
+      },
+    });
+    assert.deepStrictEqual(
+      nextNthRequest(fauxDatabase),
+      4,
+      'With override should infer value'
+    );
+  });
+});