Skip to content

feat(spanner): make affinity keys configurable via spanner_grpc_confi… #8157

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
alkatrivedi wants to merge 2 commits into
base: main
Choose a base branch
from
Open

feat(spanner): make affinity keys configurable via spanner_grpc_confi… #8157

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
eb8f9df
feat(spanner): make affinity keys configurable via spanner_grpc_confi…
alkatrivedi May 2, 2026
77436fe
gemini review comments
alkatrivedi May 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion handwritten/spanner/src/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -423,7 +423,8 @@ class Spanner extends GrpcService {
// Enable grpc-gcp support
'grpc.callInvocationTransformer': grpcGcp.gcpCallInvocationTransformer,
'grpc.channelFactoryOverride': grpcGcp.gcpChannelFactoryOverride,
'grpc.gcpApiConfig': grpcGcp.createGcpApiConfig(gcpApiConfig),
// Bypass createGcpApiConfig to preserve our custom metadata fields
'grpc.gcpApiConfig': gcpApiConfig,
Comment on lines +426 to +427
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot May 3, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Bypassing grpcGcp.createGcpApiConfig can lead to issues with the grpc-gcp library, as it typically expects a validated and processed ApiConfig object to correctly initialize channel management and affinity. Since the custom metadata fields added to the config are read directly from the JSON file in transaction.ts, they do not need to be preserved within the grpc.gcpApiConfig channel option. Bypassing this helper may inadvertently break standard affinity behavior (e.g., for regular sessions) if the library fails to parse the raw object.

Suggested change
// Bypass createGcpApiConfig to preserve our custom metadata fields
'grpc.gcpApiConfig': gcpApiConfig,
'grpc.gcpApiConfig': grpcGcp.createGcpApiConfig(gcpApiConfig),

Copy link
Copy Markdown
Contributor Author

@alkatrivedi alkatrivedi May 3, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

grpcGcp.createGcpApiConfig(gcpApiConfig) will drop our custom configuration fields (like unbindMetadataKey or affinityKeyLocation). This is because createGcpApiConfig converts the raw JSON object into a strictly typed Protobuf message instance based on the grpc_gcp.proto schema. Any fields present in the JSON that are not explicitly defined in the proto file will be filtered out and ignored." Hence, this bypassing is required

grpc,
},
options || {},
Expand Down
33 changes: 26 additions & 7 deletions handwritten/spanner/src/spanner_grpc_config.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,49 +29,67 @@
"name": ["/google.spanner.v1.Spanner/ExecuteSql"],
"affinity": {
"command": "BOUND",
"affinityKey": "session"
"affinityKey": "session",
"affinityKeyLocation": "METADATA",
"metadataKey": "x-grpc-gcp-affinity-key"
}
},
{
"name": ["/google.spanner.v1.Spanner/ExecuteStreamingSql"],
"affinity": {
"command": "BOUND",
"affinityKey": "session"
"affinityKey": "session",
"affinityKeyLocation": "METADATA",
"metadataKey": "x-grpc-gcp-affinity-key"
}
},
{
"name": ["/google.spanner.v1.Spanner/Read"],
"affinity": {
"command": "BOUND",
"affinityKey": "session"
"affinityKey": "session",
"affinityKeyLocation": "METADATA",
"metadataKey": "x-grpc-gcp-affinity-key"
}
},
{
"name": ["/google.spanner.v1.Spanner/StreamingRead"],
"affinity": {
"command": "BOUND",
"affinityKey": "session"
"affinityKey": "session",
"affinityKeyLocation": "METADATA",
"metadataKey": "x-grpc-gcp-affinity-key"
}
},
{
"name": ["/google.spanner.v1.Spanner/BeginTransaction"],
"affinity": {
"command": "BOUND",
"affinityKey": "session"
"affinityKey": "session",
"affinityKeyLocation": "METADATA",
"metadataKey": "x-grpc-gcp-affinity-key"
}
},
{
"name": ["/google.spanner.v1.Spanner/Commit"],
"affinity": {
"command": "BOUND",
"affinityKey": "session"
"affinityKey": "session",
"affinityKeyLocation": "METADATA",
"metadataKey": "x-grpc-gcp-affinity-key",
"unbindKeyLocation": "METADATA",
"unbindMetadataKey": "x-grpc-gcp-unbind"
}
},
{
"name": ["/google.spanner.v1.Spanner/Rollback"],
"affinity": {
"command": "BOUND",
"affinityKey": "session"
"affinityKey": "session",
"affinityKeyLocation": "METADATA",
"metadataKey": "x-grpc-gcp-affinity-key",
"unbindKeyLocation": "METADATA",
"unbindMetadataKey": "x-grpc-gcp-unbind"
}
},
{
Expand All @@ -90,3 +108,4 @@
}
]
}

92 changes: 88 additions & 4 deletions handwritten/spanner/src/transaction.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,21 @@ import {
} from './instrument';
import {RunTransactionOptions} from './transaction-runner';
import {injectRequestIDIntoHeaders, nextNthRequest} from './request_id_header';
import * as uuid from 'uuid';

const gcpApiConfig = require('./spanner_grpc_config.json');

// Pre-compute a map for O(1) affinity lookups
const methodToAffinityMap = new Map<string, any>();
if (gcpApiConfig && gcpApiConfig.method) {
gcpApiConfig.method.forEach((m: any) => {
if (m.name && m.affinity) {
m.name.forEach((name: string) => {
methodToAffinityMap.set(name, m.affinity);
});
}
});
}

export type Rows = Array<Row | Json>;
const RETRY_INFO_TYPE = 'type.googleapis.com/google.rpc.retryinfo';
Expand Down Expand Up @@ -296,6 +311,7 @@ export class Snapshot extends EventEmitter {
| undefined
| null;
id?: Uint8Array | string;
public _affinityKey?: string;
multiplexedSessionPreviousTransactionId?: Uint8Array | string;
ended: boolean;
metadata?: spannerClient.spanner.v1.ITransaction;
Expand Down Expand Up @@ -365,8 +381,48 @@ export class Snapshot extends EventEmitter {
this.ended = false;
this.session = session;
this.queryOptions = Object.assign({}, queryOptions);
this.request = session.request.bind(session);
this.requestStream = session.requestStream.bind(session);
// If the session is multiplexed, generate a unique affinity key (UUID) for this
// specific transaction/snapshot. This allows requests using the same shared
// multiplexed session to be distributed across different gRPC channels.
if (session.metadata && session.metadata.multiplexed) {
this._affinityKey = uuid.v4();
}
const getMetadataHeaderName = (rpcMethodName: string): string => {
const method =
rpcMethodName.charAt(0).toUpperCase() + rpcMethodName.slice(1);
const fullRpcPath = `/google.spanner.v1.Spanner/${method}`;

const affinity = methodToAffinityMap.get(fullRpcPath);
return affinity?.metadataKey || 'x-grpc-gcp-affinity-key';
};

this.request = (config: any, callback: Function) => {
if (this._affinityKey) {
const headerName = getMetadataHeaderName(config.method);
config = {
...config,
headers: {
...(config.headers || {}),
[headerName]: this._affinityKey,
},
};
}
return session.request(config, callback);
};
Comment thread
alkatrivedi marked this conversation as resolved.

this.requestStream = (config: any) => {
if (this._affinityKey) {
const headerName = getMetadataHeaderName(config.method);
config = {
...config,
headers: {
...(config.headers || {}),
[headerName]: this._affinityKey,
},
};
}
return session.requestStream(config);
};
Comment thread
alkatrivedi marked this conversation as resolved.

const readOnly = Snapshot.encodeTimestampBounds(options || {});
this._options = {readOnly};
Expand Down Expand Up @@ -2453,6 +2509,20 @@ export class Transaction extends Dml {
addLeaderAwareRoutingHeader(headers);
}

// Create a copy to avoid leaking the unbind header to the global commonHeaders_.
const requestHeaders = Object.assign({}, headers);

// Signal to grpc-gcp to unbind the affinity key and clean up memory
// since this transaction is now complete.
if (this._affinityKey) {
const affinity = methodToAffinityMap.get(
'/google.spanner.v1.Spanner/Commit',
);
const unbindHeaderName =
affinity?.unbindMetadataKey || 'x-grpc-gcp-unbind';
requestHeaders[unbindHeaderName] = 'true';
}

span.addEvent('Starting Commit');

const database = this.session.parent as Database;
Expand All @@ -2463,7 +2533,7 @@ export class Transaction extends Dml {
reqOpts,
gaxOpts: gaxOpts,
headers: injectRequestIDIntoHeaders(
headers,
requestHeaders,
this.session,
nextNthRequest(database),
1,
Expand Down Expand Up @@ -2819,13 +2889,27 @@ export class Transaction extends Dml {
addLeaderAwareRoutingHeader(headers);
}

// Create a copy to avoid leaking the unbind header to the global commonHeaders_.
const requestHeaders = Object.assign({}, headers);

// Signal to grpc-gcp to unbind the affinity key and clean up memory
// since this transaction is now complete.
if (this._affinityKey) {
const affinity = methodToAffinityMap.get(
'/google.spanner.v1.Spanner/Rollback',
);
const unbindHeaderName =
affinity?.unbindMetadataKey || 'x-grpc-gcp-unbind';
requestHeaders[unbindHeaderName] = 'true';
}

this.request(
{
client: 'SpannerClient',
method: 'rollback',
reqOpts,
gaxOpts,
headers: headers,
headers: requestHeaders,
},
(err: null | ServiceError) => {
if (err) {
Expand Down
4 changes: 1 addition & 3 deletions handwritten/spanner/test/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -232,9 +232,7 @@ describe('Spanner', () => {
'grpc.callInvocationTransformer':
fakeGrpcGcp().gcpCallInvocationTransformer,
'grpc.channelFactoryOverride': fakeGrpcGcp().gcpChannelFactoryOverride,
'grpc.gcpApiConfig': {
calledWith_: apiConfig,
},
'grpc.gcpApiConfig': apiConfig,
});

it('should localize a cached gapic client map', () => {
Expand Down
14 changes: 10 additions & 4 deletions handwritten/spanner/test/multiplexed-session.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,11 @@ describe('MultiplexedSession', () => {

return Object.assign(new Session(DATABASE, name), props, {
create: sandbox.stub().resolves(),
transaction: sandbox.stub().returns(new FakeTransaction()),
transaction: sandbox.stub().callsFake(() => {
const txn = new FakeTransaction();
(txn as any)._affinityKey = 'mock-uuid';
return txn;
}),
});
};

Expand Down Expand Up @@ -185,13 +189,15 @@ describe('MultiplexedSession', () => {
});
});

it('should pass back the session and txn', done => {
const fakeTxn = new FakeTransaction() as unknown as Transaction;
it('should pass back the session and txn with affinity key', done => {
sandbox.stub(multiplexedSession, '_getSession').resolves(fakeMuxSession);
multiplexedSession.getSession((err, session, txn) => {
assert.ifError(err);
assert.strictEqual(session, fakeMuxSession);
assert.deepStrictEqual(txn, fakeTxn);
assert(txn);
assert(txn._affinityKey);
assert.strictEqual(typeof txn._affinityKey, 'string');
assert(txn._affinityKey.length > 0);
done();
});
});
Expand Down
Loading