Skip to content

Commit

Permalink
fix(SimpleIdentifyThrottler): don't sleep negative amounts
Browse files Browse the repository at this point in the history
  • Loading branch information
didinele authored and vladfrangu committed Jan 1, 2025
1 parent fdf0b84 commit a589c6d
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 12 deletions.
47 changes: 36 additions & 11 deletions packages/ws/__tests__/util/SimpleIdentifyThrottler.test.ts
Original file line number Diff line number Diff line change
@@ -1,33 +1,58 @@
// @ts-nocheck
import { setTimeout as sleep } from 'node:timers/promises';
import { expect, test, vi } from 'vitest';
import * as timers from 'node:timers/promises';
import { expect, test, vi, beforeEach, afterEach } from 'vitest';
import { SimpleIdentifyThrottler } from '../../src/index.js';

vi.mock('node:timers/promises', () => ({
setTimeout: vi.fn(),
}));

const throttler = new SimpleIdentifyThrottler(2);
let throttler: SimpleIdentifyThrottler;
const controller = new AbortController();

vi.useFakeTimers();

const NOW = vi.fn().mockReturnValue(Date.now());
const TIME = Date.now();
const NOW = vi.fn().mockReturnValue(TIME);
global.Date.now = NOW;

const sleep = vi.spyOn(timers, 'setTimeout');

beforeEach(() => {
throttler = new SimpleIdentifyThrottler(2);
});

afterEach(() => {
sleep.mockClear();
});

test('basic case', async () => {
// Those shouldn't wait since they're in different keys

await throttler.waitForIdentify(0);
await throttler.waitForIdentify(0, controller.signal);
expect(sleep).not.toHaveBeenCalled();

await throttler.waitForIdentify(1);
await throttler.waitForIdentify(1, controller.signal);
expect(sleep).not.toHaveBeenCalled();

// Those should wait

await throttler.waitForIdentify(2);
await throttler.waitForIdentify(2, controller.signal);
expect(sleep).toHaveBeenCalledTimes(1);

await throttler.waitForIdentify(3);
await throttler.waitForIdentify(3, controller.signal);
expect(sleep).toHaveBeenCalledTimes(2);
});

test('does not call sleep with a negative time', async () => {
await throttler.waitForIdentify(0, controller.signal);
expect(sleep).not.toHaveBeenCalled();

await throttler.waitForIdentify(0, controller.signal);
expect(sleep).toHaveBeenCalledTimes(1);

// By overshooting, we're simulating a bug that existed prior to this test, where-in by enough time
// passing before the shard tried to identify for a subsequent time, the passed value would end up being negative
// (and this was unchecked). Node simply treats that as 1ms, so it wasn't particularly harmful, but they
// recently introduced a warning for it, so we want to avoid that.
NOW.mockReturnValueOnce(TIME + 10_000);
await throttler.waitForIdentify(0, controller.signal);
expect(sleep).toHaveBeenCalledTimes(1);
});
2 changes: 1 addition & 1 deletion packages/ws/src/throttling/SimpleIdentifyThrottler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ export class SimpleIdentifyThrottler implements IIdentifyThrottler {

try {
const diff = state.resetsAt - Date.now();
if (diff <= 5_000) {
if (diff > 0 && diff <= 5_000) {
// To account for the latency the IDENTIFY payload goes through, we add a bit more wait time
const time = diff + Math.random() * 1_500;
await sleep(time);
Expand Down

0 comments on commit a589c6d

Please sign in to comment.