Skip to content

Commit

Permalink
feat: send retry attempt header to ease debugging (#1068)
Browse files Browse the repository at this point in the history
* feat: send retry attempt header to ease debugging

* style

* lint
  • Loading branch information
igorbernstein2 authored Apr 13, 2022
1 parent 3a02e49 commit 37f9b3c
Show file tree
Hide file tree
Showing 2 changed files with 78 additions and 2 deletions.
20 changes: 20 additions & 0 deletions src/table.ts
Original file line number Diff line number Diff line change
Expand Up @@ -743,6 +743,7 @@ Please use the format 'prezzy' or '${instance.name}/tables/prezzy'.`);
const hasLimit = rowsLimit !== 0;
let rowsRead = 0;
let numConsecutiveErrors = 0;
let numRequestsMade = 0;
let retryTimer: NodeJS.Timeout | null;

rowKeys = options.keys || [];
Expand Down Expand Up @@ -917,6 +918,11 @@ Please use the format 'prezzy' or '${instance.name}/tables/prezzy'.`);
reqOpts.rowsLimit = rowsLimit - rowsRead;
}

options.gaxOptions = populateAttemptHeader(
numRequestsMade,
options.gaxOptions
);

const requestStream = this.bigtable.request({
client: 'BigtableClient',
method: 'readRows',
Expand Down Expand Up @@ -970,6 +976,7 @@ Please use the format 'prezzy' or '${instance.name}/tables/prezzy'.`);
return;
}
numConsecutiveErrors++;
numRequestsMade++;
if (
numConsecutiveErrors <= maxRetries &&
(RETRYABLE_STATUS_CODES.has(error.code) || isRstStreamError(error))
Expand Down Expand Up @@ -1614,6 +1621,11 @@ Please use the format 'prezzy' or '${instance.name}/tables/prezzy'.`);
},
};

options.gaxOptions = populateAttemptHeader(
numRequestsMade,
options.gaxOptions
);

this.bigtable
.request<google.bigtable.v2.MutateRowsResponse>({
client: 'BigtableClient',
Expand Down Expand Up @@ -2072,6 +2084,14 @@ function getNextDelay(numConsecutiveErrors: number, config: BackoffSettings) {
return Math.min(calculatedNextRetryDelay, config.maxRetryDelayMillis);
}

function populateAttemptHeader(attempt: number, gaxOpts?: CallOptions) {
gaxOpts = gaxOpts || {};
gaxOpts.otherArgs = gaxOpts.otherArgs || {};
gaxOpts.otherArgs.headers = gaxOpts.otherArgs.headers || {};
gaxOpts.otherArgs.headers['bigtable-attempt'] = attempt;
return gaxOpts;
}

export interface GoogleInnerError {
reason?: string;
message?: string;
Expand Down
60 changes: 58 additions & 2 deletions test/table.ts
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ import {Family} from '../src/family.js';
import {Mutation} from '../src/mutation.js';
import {Row} from '../src/row.js';
import * as tblTypes from '../src/table';
import {Bigtable} from '../src';
import {Bigtable, RequestOptions} from '../src';
import {EventEmitter} from 'events';

const sandbox = sinon.createSandbox();
Expand Down Expand Up @@ -544,7 +544,9 @@ describe('Bigtable/Table', () => {
assert.strictEqual(config.method, 'readRows');
assert.strictEqual(config.reqOpts.tableName, TABLE_NAME);
assert.strictEqual(config.reqOpts.appProfileId, undefined);
assert.strictEqual(config.gaxOpts, undefined);
assert.deepStrictEqual(config.gaxOpts, {
otherArgs: {headers: {'bigtable-attempt': 0}},
});
done();
};
table.createReadStream();
Expand Down Expand Up @@ -2570,6 +2572,7 @@ describe('Bigtable/Table', () => {
let fakeStatuses: any;
// eslint-disable-next-line @typescript-eslint/no-explicit-any
let entryRequests: any;
const requestArgs: RequestOptions[] = [];

beforeEach(() => {
entryRequests = [];
Expand Down Expand Up @@ -2599,6 +2602,7 @@ describe('Bigtable/Table', () => {
];
// eslint-disable-next-line @typescript-eslint/no-explicit-any
table.bigtable.request = (config: any) => {
requestArgs.push(JSON.parse(JSON.stringify(config)));
entryRequests.push(config.reqOpts.entries);
const stream = new PassThrough({
objectMode: true,
Expand All @@ -2612,6 +2616,25 @@ describe('Bigtable/Table', () => {
};
});

it('should send attempt header', done => {
table.mutate(entries, () => {
assert.strictEqual(requestArgs.length, 2);
assert.strictEqual(
(requestArgs[0].gaxOpts as any)['otherArgs']['headers'][
'bigtable-attempt'
],
0
);
assert.strictEqual(
(requestArgs[1].gaxOpts as any)['otherArgs']['headers'][
'bigtable-attempt'
],
1
);
done();
});
});

it('should succeed after a retry', done => {
table.maxRetries = 1;
table.mutate(entries, done);
Expand All @@ -2630,17 +2653,20 @@ describe('Bigtable/Table', () => {

describe('rpc level retries', () => {
let emitters: EventEmitter[] | null; // = [((stream: Writable) => { stream.push([{ key: 'a' }]);
let requestArgs: RequestOptions[] = [];

// eslint-disable-next-line @typescript-eslint/no-explicit-any
let entryRequests: any;

beforeEach(() => {
emitters = null; // This needs to be assigned in each test case.

requestArgs = [];
entryRequests = [];

// eslint-disable-next-line @typescript-eslint/no-explicit-any
table.bigtable.request = (config: any) => {
requestArgs.push(JSON.parse(JSON.stringify(config)));
entryRequests.push(config.reqOpts.entries);
const stream = new PassThrough({
objectMode: true,
Expand Down Expand Up @@ -2707,6 +2733,36 @@ describe('Bigtable/Table', () => {
done();
});
});

it('should send attempt header', done => {
const error = new Error('retryable') as ServiceError;
error.code = 14; // Unavailable
emitters = [
((stream: Writable) => {
stream.emit('error', error);
}) as {} as EventEmitter,
((stream: Writable) => {
stream.end();
}) as {} as EventEmitter,
];
table.maxRetries = 1;
table.mutate(entries, () => {
assert.strictEqual(requestArgs.length, 2);
assert.strictEqual(
(requestArgs[0].gaxOpts as any)['otherArgs']['headers'][
'bigtable-attempt'
],
0
);
assert.strictEqual(
(requestArgs[1].gaxOpts as any)['otherArgs']['headers'][
'bigtable-attempt'
],
1
);
done();
});
});
});
});

Expand Down

0 comments on commit 37f9b3c

Please sign in to comment.