diff --git a/handwritten/pubsub/test/publisher/message-queues.ts b/handwritten/pubsub/test/publisher/message-queues.ts index c6d0b330998..31d7e3c3a72 100644 --- a/handwritten/pubsub/test/publisher/message-queues.ts +++ b/handwritten/pubsub/test/publisher/message-queues.ts @@ -14,7 +14,7 @@ * limitations under the License. */ -import {loggingUtils, ServiceError} from 'google-gax'; +import {ServiceError} from 'google-gax'; import * as assert from 'assert'; import {describe, it, before, beforeEach, afterEach} from 'mocha'; import {EventEmitter} from 'events'; @@ -26,10 +26,22 @@ import * as p from '../../src/publisher'; import * as b from '../../src/publisher/message-batch'; import * as q from '../../src/publisher/message-queues'; import {PublishError} from '../../src/publisher/publish-error'; +import {Duration} from '../../src/temporal'; +import * as otel from '../../src/telemetry-tracing'; +import {exporter} from '../tracing'; import {FakeLog, TestUtils} from '../test-utils'; +import {google} from '../../protos/protos'; + +class FakePubSub { + isIdResolved = true; + async getClientConfig(): Promise { + return {}; + } +} class FakeTopic { name = 'projects/foo/topics/fake-topic'; + pubsub = new FakePubSub(); // eslint-disable-next-line @typescript-eslint/no-unused-vars request(config: RequestConfig, callback: RequestCallback): void {} } @@ -40,6 +52,9 @@ class FakePublisher { topic: FakeTopic; settings: p.PublishOptions; flowControl: FakeFlowControl; + hedgingOptions?: p.ResolvedHedgingOptions; + private tokenBucket_?: p.HedgingTokenBucket; + hedgingScheduler?: p.HedgingScheduler; constructor(topic: FakeTopic) { this.topic = topic; this.settings = { @@ -47,6 +62,29 @@ class FakePublisher { }; this.flowControl = new FakeFlowControl(); } + get tokenBucket(): p.HedgingTokenBucket | undefined { + return this.tokenBucket_; + } + enableHedging(options: p.HedgingOptions = {}): void { + const resolved = p.HedgingTokenBucket.validateAndResolveHedgingOptions( + options, + this.settings.messageOrdering, + this.settings.gaxOpts, + ); + this.hedgingOptions = resolved; + if (resolved) { + const bucket = new p.HedgingTokenBucket(resolved); + this.tokenBucket_ = bucket; + this.hedgingScheduler = new p.HedgingScheduler(resolved, bucket, () => { + q.logs.publishHedged.debug( + 'Hedging rate limited due to lack of tokens.', + ); + }); + } + } + refillTokenBucket(): void { + this.tokenBucket_?.refillTokenBucket(); + } } class FakeMessageBatch { @@ -187,10 +225,7 @@ describe('Message Queues', () => { fakeLog.remove(); assert.strictEqual(fakeLog.called, true); - assert.strictEqual( - fakeLog.fields!.severity, - 'INFO', - ); + assert.strictEqual(fakeLog.fields!.severity, 'INFO'); assert.strictEqual(fakeLog.args![1] as string, 'test'); }); @@ -377,7 +412,7 @@ describe('Message Queues', () => { spies = [sandbox.spy(), sandbox.spy()] as p.PublishCallback[]; }); - it('should begin another publish(drain) if there are pending batches', done => { + it('should begin another publish(drain) if there are pending batches', async () => { const stub = sandbox.stub(queue, '_publish'); let once = false; stub.callsFake(async () => { @@ -394,12 +429,8 @@ describe('Message Queues', () => { queue.batch = new FakeMessageBatch(); queue.batch.messages = fakeMessages; queue.batch.callbacks = spies; - void queue.publishDrain().then(() => { - process.nextTick(() => { - assert.strictEqual(stub.callCount, 2); - done(); - }); - }); + await queue.publishDrain(); + assert.strictEqual(stub.callCount, 2); }); it('should not begin another publish(non-drain) if there are pending batches', async () => { @@ -424,17 +455,13 @@ describe('Message Queues', () => { assert.strictEqual(stub.callCount, 1); }); - it('should emit "drain" if there is nothing left to publish', done => { + it('should emit "drain" if there is nothing left to publish', async () => { const spy = sandbox.spy(); sandbox.stub(queue, '_publish').callsFake(async () => {}); queue.on('drain', spy); - void queue.publish('test').then(() => { - process.nextTick(() => { - assert.strictEqual(spy.callCount, 1); - done(); - }); - }); + await queue.publish('test'); + assert.strictEqual(spy.callCount, 1); }); }); }); @@ -804,4 +831,413 @@ describe('Message Queues', () => { }); }); }); + + describe('Publish Hedging in Queues', () => { + interface CapturedPublishCall { + config: RequestConfig; + callback: RequestCallback; + } + + let clock: sinon.SinonFakeTimers; + let fakePub: FakePublisher; + let queue: q.Queue; + + beforeEach(() => { + clock = TestUtils.useFakeTimers(sandbox, 10000); + fakePub = publisher as unknown as FakePublisher; + queue = new Queue(publisher); + exporter.reset(); + }); + + afterEach(() => { + fakePub.hedgingScheduler?.clear(); + clock.restore(); + exporter.reset(); + otel.setGloballyEnabled(false); + }); + + it('triggers a hedged publish with default settings once tokens are available', async () => { + fakePub.enableHedging({}); + for (let i = 0; i < 10; i++) { + fakePub.refillTokenBucket(); + } + assert.strictEqual(fakePub.tokenBucket?.getTokenBalance(), 1.0); + + const calls: CapturedPublishCall[] = []; + sandbox + .stub(topic, 'request') + .callsFake( + ( + config: RequestConfig, + callback: RequestCallback, + ) => { + calls.push({config, callback}); + }, + ); + + const cbSpy = sandbox.spy(); + const publishPromise = queue._publish([{}], [cbSpy], 0, 'test'); + + // Initial attempt is sent immediately + assert.strictEqual(calls.length, 1); + assert.ok(calls[0].config.gaxOpts?.signal); + + // Advance 1000ms (default hedgeDelay) -> hedged attempt 1 is sent + clock.tick(1000); + assert.strictEqual(calls.length, 2); + assert.strictEqual(calls[1].config.gaxOpts?.retry, null); + assert.strictEqual(calls[1].config.gaxOpts?.timeout, 10000); + + // Complete hedged attempt -> refills 0.1 token + calls[1].callback(null, {messageIds: ['hedged-msg-1']}); + await publishPromise; + + assert.strictEqual( + cbSpy.calledOnceWithExactly(null, 'hedged-msg-1'), + true, + ); + assert.strictEqual(calls[0].config.gaxOpts?.signal?.aborted, true); + assert.strictEqual(fakePub.tokenBucket?.getTokenBalance(), 0.1); + }); + + it('cancels hedged attempt and emits standard publish end event when original attempt succeeds after hedging', async () => { + otel.setGloballyEnabled(true); + fakePub.enableHedging({ + hedgeDelay: Duration.from({milliseconds: 300}), + maxTokens: 10, + refillRatio: 0.2, + }); + for (let i = 0; i < 5; i++) { + fakePub.refillTokenBucket(); + } + + const calls: CapturedPublishCall[] = []; + sandbox + .stub(topic, 'request') + .callsFake( + ( + config: RequestConfig, + callback: RequestCallback, + ) => { + calls.push({config, callback}); + }, + ); + + const message: p.PubsubMessage = {data: Buffer.from('hello')}; + const parentSpan = otel.PubsubSpans.createPublisherSpan( + message, + topic.name, + 'test', + ); + message.parentSpan = parentSpan; + + const cbSpy = sandbox.spy(); + const publishPromise = queue._publish([message], [cbSpy], 5, 'test'); + + clock.tick(300); + assert.strictEqual(calls.length, 2); + + // Original attempt (index 0) finishes first + calls[0].callback(null, {messageIds: ['orig-id']}); + await publishPromise; + + assert.strictEqual(cbSpy.calledOnceWithExactly(null, 'orig-id'), true); + assert.strictEqual(calls[0].config.gaxOpts?.signal?.aborted, false); + assert.strictEqual(calls[1].config.gaxOpts?.signal?.aborted, true); + + const createSpan = exporter + .getFinishedSpans() + .find(s => s.name.endsWith('create')); + assert.ok(createSpan); + const eventNames = createSpan.events.map(e => e.name); + assert.deepStrictEqual(eventNames, [ + 'publish start', + 'publish start (hedged)', + 'publish end', + ]); + }); + + it('cancels original attempt and emits hedged publish end event and debug log when hedged attempt succeeds first', async () => { + otel.setGloballyEnabled(true); + fakePub.enableHedging({ + hedgeDelay: Duration.from({milliseconds: 300}), + maxTokens: 10, + refillRatio: 0.2, + }); + for (let i = 0; i < 5; i++) { + fakePub.refillTokenBucket(); + } + + const calls: CapturedPublishCall[] = []; + sandbox + .stub(topic, 'request') + .callsFake( + ( + config: RequestConfig, + callback: RequestCallback, + ) => { + calls.push({config, callback}); + }, + ); + + const message: p.PubsubMessage = {data: Buffer.from('hello')}; + const parentSpan = otel.PubsubSpans.createPublisherSpan( + message, + topic.name, + 'test', + ); + message.parentSpan = parentSpan; + + const fakeLog = new FakeLog(q.logs.publishHedged); + const cbSpy = sandbox.spy(); + const publishPromise = queue._publish([message], [cbSpy], 5, 'test'); + + clock.tick(300); + assert.strictEqual(calls.length, 2); + fakeLog.remove(); + + assert.strictEqual(fakeLog.called, true); + assert.strictEqual(fakeLog.fields?.severity, 'DEBUG'); + assert.strictEqual(fakeLog.args?.[1], 1); + + // Hedged attempt (index 1) finishes first + calls[1].callback(null, {messageIds: ['hedged-id']}); + await publishPromise; + + assert.strictEqual(cbSpy.calledOnceWithExactly(null, 'hedged-id'), true); + assert.strictEqual(calls[0].config.gaxOpts?.signal?.aborted, true); + assert.strictEqual(calls[1].config.gaxOpts?.signal?.aborted, false); + + const createSpan = exporter + .getFinishedSpans() + .find(s => s.name.endsWith('create')); + assert.ok(createSpan); + const eventNames = createSpan.events.map(e => e.name); + assert.deepStrictEqual(eventNames, [ + 'publish start', + 'publish start (hedged)', + 'publish end (hedged)', + ]); + }); + + it('discards hedged request error and resolves when original request succeeds', async () => { + fakePub.enableHedging({ + hedgeDelay: Duration.from({milliseconds: 250}), + maxTokens: 10, + refillRatio: 0.2, + }); + for (let i = 0; i < 5; i++) { + fakePub.refillTokenBucket(); + } + + const calls: CapturedPublishCall[] = []; + sandbox + .stub(topic, 'request') + .callsFake( + ( + config: RequestConfig, + callback: RequestCallback, + ) => { + calls.push({config, callback}); + }, + ); + + const cbSpy = sandbox.spy(); + const publishPromise = queue._publish([{}], [cbSpy], 0, 'test'); + + clock.tick(250); + assert.strictEqual(calls.length, 2); + + // Hedged attempt fails while original attempt is still in flight + const hedgedError = new Error('hedged transient error') as ServiceError; + calls[1].callback(hedgedError); + await Promise.resolve(); + assert.strictEqual(cbSpy.called, false); + + // Original attempt succeeds + calls[0].callback(null, {messageIds: ['orig-after-hedge-err']}); + await publishPromise; + + assert.strictEqual( + cbSpy.calledOnceWithExactly(null, 'orig-after-hedge-err'), + true, + ); + }); + + it('cancels hedged attempt and rejects when original request fails permanently', async () => { + fakePub.enableHedging({ + hedgeDelay: Duration.from({milliseconds: 250}), + maxTokens: 10, + refillRatio: 0.2, + }); + for (let i = 0; i < 5; i++) { + fakePub.refillTokenBucket(); + } + + const calls: CapturedPublishCall[] = []; + sandbox + .stub(topic, 'request') + .callsFake( + ( + config: RequestConfig, + callback: RequestCallback, + ) => { + calls.push({config, callback}); + }, + ); + + const cbSpy = sandbox.spy(); + const publishPromise = queue._publish([{}], [cbSpy], 0, 'test'); + + clock.tick(250); + assert.strictEqual(calls.length, 2); + + const permanentError = new Error('permanent failure') as ServiceError; + calls[0].callback(permanentError); + + await assert.rejects(publishPromise, permanentError); + assert.strictEqual(cbSpy.calledOnceWithExactly(permanentError), true); + assert.strictEqual(calls[1].config.gaxOpts?.signal?.aborted, true); + }); + + it('sends multiple hedged attempts at hedgeDelay intervals and cancels all others when attempt 2 wins', async () => { + fakePub.enableHedging({ + hedgeDelay: Duration.from({milliseconds: 200}), + maxTokens: 10, + refillRatio: 0.2, + }); + for (let i = 0; i < 10; i++) { + fakePub.refillTokenBucket(); + } + assert.strictEqual(fakePub.tokenBucket?.getTokenBalance(), 2.0); + + const calls: CapturedPublishCall[] = []; + sandbox + .stub(topic, 'request') + .callsFake( + ( + config: RequestConfig, + callback: RequestCallback, + ) => { + calls.push({config, callback}); + }, + ); + + const cbSpy = sandbox.spy(); + const publishPromise = queue._publish([{}], [cbSpy], 0, 'test'); + + assert.strictEqual(calls.length, 1); + + clock.tick(200); + assert.strictEqual(calls.length, 2); + + clock.tick(200); + assert.strictEqual(calls.length, 3); + + // Second hedged attempt (calls[2]) succeeds + calls[2].callback(null, {messageIds: ['hedge-2-id']}); + await publishPromise; + + assert.strictEqual(cbSpy.calledOnceWithExactly(null, 'hedge-2-id'), true); + assert.strictEqual(calls[0].config.gaxOpts?.signal?.aborted, true); + assert.strictEqual(calls[1].config.gaxOpts?.signal?.aborted, true); + assert.strictEqual(calls[2].config.gaxOpts?.signal?.aborted, false); + }); + + it('logs rate limiting when token bucket is empty and resumes hedging after refill', async () => { + fakePub.enableHedging({ + hedgeDelay: Duration.from({milliseconds: 200}), + maxTokens: 5, + refillRatio: 0.2, + }); + assert.strictEqual(fakePub.tokenBucket?.getTokenBalance(), 0); + + const calls: CapturedPublishCall[] = []; + sandbox + .stub(topic, 'request') + .callsFake( + ( + config: RequestConfig, + callback: RequestCallback, + ) => { + calls.push({config, callback}); + }, + ); + + const fakeLog = new FakeLog(q.logs.publishHedged); + const firstPromise = queue._publish([{}], [sandbox.spy()], 0, 'test'); + + clock.tick(200); + fakeLog.remove(); + + // Bucket is empty -> no hedged call was made, rate limit debug log emitted + assert.strictEqual(calls.length, 1); + assert.strictEqual(fakeLog.called, true); + assert.strictEqual(fakeLog.fields?.severity, 'DEBUG'); + assert.strictEqual( + fakeLog.args?.[0], + 'Hedging rate limited due to lack of tokens.', + ); + + calls[0].callback(null, {messageIds: ['first']}); + await firstPromise; + assert.strictEqual(fakePub.tokenBucket?.getTokenBalance(), 0.2); + + // Complete 4 more fast batches to reach 1.0 token + for (let i = 0; i < 4; i++) { + const pBatch = queue._publish([{}], [sandbox.spy()], 0, 'test'); + calls[calls.length - 1].callback(null, {messageIds: [`fast-${i}`]}); + await pBatch; + } + assert.strictEqual(fakePub.tokenBucket?.getTokenBalance(), 1.0); + + // Next slow batch should now hedge! + const callsBeforeSlow = calls.length; + const slowPromise = queue._publish([{}], [sandbox.spy()], 0, 'test'); + clock.tick(200); + assert.strictEqual(calls.length, callsBeforeSlow + 2); + + calls[calls.length - 1].callback(null, { + messageIds: ['hedged-after-refill'], + }); + await slowPromise; + }); + + it('never hedges on OrderedQueue even if hedgingScheduler is present on publisher', async () => { + fakePub.enableHedging({ + hedgeDelay: Duration.from({milliseconds: 200}), + maxTokens: 10, + refillRatio: 0.2, + }); + for (let i = 0; i < 5; i++) { + fakePub.refillTokenBucket(); + } + + const orderedQueue = new OrderedQueue(publisher, 'order-key'); + const calls: CapturedPublishCall[] = []; + sandbox + .stub(topic, 'request') + .callsFake( + ( + config: RequestConfig, + callback: RequestCallback, + ) => { + calls.push({config, callback}); + }, + ); + + const publishPromise = orderedQueue._publish( + [{orderingKey: 'order-key'}], + [sandbox.spy()], + 0, + 'test', + ); + + clock.tick(500); + assert.strictEqual(calls.length, 1); + + calls[0].callback(null, {messageIds: ['ordered-1']}); + await publishPromise; + }); + }); }); diff --git a/handwritten/pubsub/test/pubsub.ts b/handwritten/pubsub/test/pubsub.ts index b2ecbfa7d29..0427c69dbf8 100644 --- a/handwritten/pubsub/test/pubsub.ts +++ b/handwritten/pubsub/test/pubsub.ts @@ -196,7 +196,7 @@ describe('PubSub', () => { './snapshot': {Snapshot: FakeSnapshot}, './subscription': {Subscription}, './topic': {Topic: FakeTopic}, - './index': {v1:v1Override}, + './index': {v1: v1Override}, './util': fakeUtil, }).PubSub; }); @@ -1676,6 +1676,59 @@ describe('PubSub', () => { pubsub.request(CONFIG, assert.ifError); }); + + it('should immediately invoke callback with CANCELLED error if gaxOpts.signal is already aborted', done => { + const abortController = new AbortController(); + abortController.abort(); + + let getClientCalled = false; + pubsub.getClient_ = () => { + getClientCalled = true; + }; + + const configWithSignal: pubsubTypes.RequestConfig = { + client: 'PublisherClient', + method: 'publish', + reqOpts: {}, + gaxOpts: {signal: abortController.signal}, + }; + + pubsub.request(configWithSignal, (err: gax.ServiceError | null) => { + assert.strictEqual(getClientCalled, false); + assert.ok(err); + assert.strictEqual(err.code, gax.Status.CANCELLED); + done(); + }); + }); + + it('should invoke callback with CANCELLED error if gaxOpts.signal is aborted while getClient_ is in flight', done => { + const abortController = new AbortController(); + let clientMethodCalled = false; + const fakeClient = { + publish() { + clientMethodCalled = true; + }, + }; + + pubsub.getClient_ = (_config, callback) => { + abortController.abort(); + callback(null, fakeClient as unknown as gax.ClientStub); + }; + + const configWithSignal: pubsubTypes.RequestConfig = { + client: 'PublisherClient', + method: 'publish', + reqOpts: {}, + gaxOpts: {signal: abortController.signal}, + }; + + pubsub.request(configWithSignal, (err: gax.ServiceError | null) => { + assert.strictEqual(clientMethodCalled, false); + assert.ok(err); + assert.strictEqual(err.code, gax.Status.CANCELLED); + done(); + }); + }); }); describe('snapshot', () => {