mirror of
https://github.com/discordjs/discord.js.git
synced 2024-08-21 22:44:43 +12:00
refactor(WebSocketManager): passing in strategy (#9122)
* refactor(WebSocketManager): passing in strategy * chore: update tests * chore: requested nits --------- Co-authored-by: kodiakhq[bot] <49736102+kodiakhq[bot]@users.noreply.github.com>
This commit is contained in:
parent
c6f9c50ba9
commit
5c5a5832b9
5 changed files with 39 additions and 26 deletions
|
@ -107,12 +107,11 @@ const manager = new WebSocketManager({
|
||||||
intents: 0,
|
intents: 0,
|
||||||
rest,
|
rest,
|
||||||
shardCount: 6,
|
shardCount: 6,
|
||||||
|
// This will cause 3 workers to spawn, 2 shards per each
|
||||||
|
buildStrategy: (manager) => new WorkerShardingStrategy(manager, { shardsPerWorker: 2 }),
|
||||||
|
// Or maybe you want all your shards under a single worker
|
||||||
|
buildStrategy: (manager) => new WorkerShardingStrategy(manager, { shardsPerWorker: 'all' }),
|
||||||
});
|
});
|
||||||
|
|
||||||
// This will cause 3 workers to spawn, 2 shards per each
|
|
||||||
manager.setStrategy(new WorkerShardingStrategy(manager, { shardsPerWorker: 2 }));
|
|
||||||
// Or maybe you want all your shards under a single worker
|
|
||||||
manager.setStrategy(new WorkerShardingStrategy(manager, { shardsPerWorker: 'all' }));
|
|
||||||
```
|
```
|
||||||
|
|
||||||
**Note**: By default, this will cause the workers to effectively only be responsible for the WebSocket connection, they simply pass up all the events back to the main process for the manager to emit. If you want to have the workers handle events as well, you can pass in a `workerPath` option to the `WorkerShardingStrategy` constructor:
|
**Note**: By default, this will cause the workers to effectively only be responsible for the WebSocket connection, they simply pass up all the events back to the main process for the manager to emit. If you want to have the workers handle events as well, you can pass in a `workerPath` option to the `WorkerShardingStrategy` constructor:
|
||||||
|
@ -126,14 +125,12 @@ const manager = new WebSocketManager({
|
||||||
token: process.env.DISCORD_TOKEN,
|
token: process.env.DISCORD_TOKEN,
|
||||||
intents: 0,
|
intents: 0,
|
||||||
rest,
|
rest,
|
||||||
});
|
buildStrategy: (manager) =>
|
||||||
|
|
||||||
manager.setStrategy(
|
|
||||||
new WorkerShardingStrategy(manager, {
|
new WorkerShardingStrategy(manager, {
|
||||||
shardsPerWorker: 2,
|
shardsPerWorker: 2,
|
||||||
workerPath: './worker.js',
|
workerPath: './worker.js',
|
||||||
}),
|
}),
|
||||||
);
|
});
|
||||||
```
|
```
|
||||||
|
|
||||||
And your `worker.ts` file:
|
And your `worker.ts` file:
|
||||||
|
|
|
@ -163,6 +163,7 @@ test('spawn, connect, send a message, session info, and destroy', async () => {
|
||||||
shardIds: [0, 1],
|
shardIds: [0, 1],
|
||||||
retrieveSessionInfo: mockRetrieveSessionInfo,
|
retrieveSessionInfo: mockRetrieveSessionInfo,
|
||||||
updateSessionInfo: mockUpdateSessionInfo,
|
updateSessionInfo: mockUpdateSessionInfo,
|
||||||
|
buildStrategy: (manager) => new WorkerShardingStrategy(manager, { shardsPerWorker: 'all' }),
|
||||||
});
|
});
|
||||||
|
|
||||||
const managerEmitSpy = vi.spyOn(manager, 'emit');
|
const managerEmitSpy = vi.spyOn(manager, 'emit');
|
||||||
|
@ -191,9 +192,6 @@ test('spawn, connect, send a message, session info, and destroy', async () => {
|
||||||
},
|
},
|
||||||
}));
|
}));
|
||||||
|
|
||||||
const strategy = new WorkerShardingStrategy(manager, { shardsPerWorker: 'all' });
|
|
||||||
manager.setStrategy(strategy);
|
|
||||||
|
|
||||||
await manager.connect();
|
await manager.connect();
|
||||||
expect(mockConstructor).toHaveBeenCalledWith(
|
expect(mockConstructor).toHaveBeenCalledWith(
|
||||||
expect.stringContaining('defaultWorker.js'),
|
expect.stringContaining('defaultWorker.js'),
|
||||||
|
|
|
@ -177,14 +177,21 @@ test('strategies', async () => {
|
||||||
public destroy = vi.fn();
|
public destroy = vi.fn();
|
||||||
|
|
||||||
public send = vi.fn();
|
public send = vi.fn();
|
||||||
|
|
||||||
|
public fetchStatus = vi.fn();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
const strategy = new MockStrategy();
|
||||||
|
|
||||||
const rest = new REST().setAgent(mockAgent).setToken('A-Very-Fake-Token');
|
const rest = new REST().setAgent(mockAgent).setToken('A-Very-Fake-Token');
|
||||||
const shardIds = [0, 1, 2];
|
const shardIds = [0, 1, 2];
|
||||||
const manager = new WebSocketManager({ token: 'A-Very-Fake-Token', intents: 0, rest, shardIds });
|
const manager = new WebSocketManager({
|
||||||
|
token: 'A-Very-Fake-Token',
|
||||||
const strategy = new MockStrategy();
|
intents: 0,
|
||||||
manager.setStrategy(strategy);
|
rest,
|
||||||
|
shardIds,
|
||||||
|
buildStrategy: () => strategy,
|
||||||
|
});
|
||||||
|
|
||||||
const data: APIGatewayBotInfo = {
|
const data: APIGatewayBotInfo = {
|
||||||
shards: 1,
|
shards: 1,
|
||||||
|
|
|
@ -2,6 +2,7 @@ import process from 'node:process';
|
||||||
import { Collection } from '@discordjs/collection';
|
import { Collection } from '@discordjs/collection';
|
||||||
import { lazy } from '@discordjs/util';
|
import { lazy } from '@discordjs/util';
|
||||||
import { APIVersion, GatewayOpcodes } from 'discord-api-types/v10';
|
import { APIVersion, GatewayOpcodes } from 'discord-api-types/v10';
|
||||||
|
import { SimpleShardingStrategy } from '../strategies/sharding/SimpleShardingStrategy.js';
|
||||||
import type { SessionInfo, OptionalWebSocketManagerOptions } from '../ws/WebSocketManager.js';
|
import type { SessionInfo, OptionalWebSocketManagerOptions } from '../ws/WebSocketManager.js';
|
||||||
import type { SendRateLimitState } from '../ws/WebSocketShard.js';
|
import type { SendRateLimitState } from '../ws/WebSocketShard.js';
|
||||||
|
|
||||||
|
@ -27,6 +28,7 @@ const getDefaultSessionStore = lazy(() => new Collection<number, SessionInfo | n
|
||||||
* Default options used by the manager
|
* Default options used by the manager
|
||||||
*/
|
*/
|
||||||
export const DefaultWebSocketManagerOptions = {
|
export const DefaultWebSocketManagerOptions = {
|
||||||
|
buildStrategy: (manager) => new SimpleShardingStrategy(manager),
|
||||||
shardCount: null,
|
shardCount: null,
|
||||||
shardIds: null,
|
shardIds: null,
|
||||||
largeThreshold: null,
|
largeThreshold: null,
|
||||||
|
|
|
@ -11,7 +11,6 @@ import {
|
||||||
type GatewaySendPayload,
|
type GatewaySendPayload,
|
||||||
} from 'discord-api-types/v10';
|
} from 'discord-api-types/v10';
|
||||||
import type { IShardingStrategy } from '../strategies/sharding/IShardingStrategy';
|
import type { IShardingStrategy } from '../strategies/sharding/IShardingStrategy';
|
||||||
import { SimpleShardingStrategy } from '../strategies/sharding/SimpleShardingStrategy.js';
|
|
||||||
import { DefaultWebSocketManagerOptions, type CompressionMethod, type Encoding } from '../utils/constants.js';
|
import { DefaultWebSocketManagerOptions, type CompressionMethod, type Encoding } from '../utils/constants.js';
|
||||||
import type { WebSocketShardDestroyOptions, WebSocketShardEventsMap } from './WebSocketShard.js';
|
import type { WebSocketShardDestroyOptions, WebSocketShardEventsMap } from './WebSocketShard.js';
|
||||||
|
|
||||||
|
@ -71,6 +70,20 @@ export interface RequiredWebSocketManagerOptions {
|
||||||
* Optional additional configuration for the WebSocketManager
|
* Optional additional configuration for the WebSocketManager
|
||||||
*/
|
*/
|
||||||
export interface OptionalWebSocketManagerOptions {
|
export interface OptionalWebSocketManagerOptions {
|
||||||
|
/**
|
||||||
|
* Builds the strategy to use for sharding
|
||||||
|
*
|
||||||
|
* @example
|
||||||
|
* ```ts
|
||||||
|
* const manager = new WebSocketManager({
|
||||||
|
* token: process.env.DISCORD_TOKEN,
|
||||||
|
* intents: 0, // for no intents
|
||||||
|
* rest,
|
||||||
|
* buildStrategy: (manager) => new WorkerShardingStrategy(manager, { shardsPerWorker: 2 }),
|
||||||
|
* });
|
||||||
|
* ```
|
||||||
|
*/
|
||||||
|
buildStrategy(manager: WebSocketManager): IShardingStrategy;
|
||||||
/**
|
/**
|
||||||
* The compression method to use
|
* The compression method to use
|
||||||
*
|
*
|
||||||
|
@ -192,16 +205,12 @@ export class WebSocketManager extends AsyncEventEmitter<ManagerShardEventsMap> {
|
||||||
*
|
*
|
||||||
* @defaultValue `SimpleShardingStrategy`
|
* @defaultValue `SimpleShardingStrategy`
|
||||||
*/
|
*/
|
||||||
private strategy: IShardingStrategy = new SimpleShardingStrategy(this);
|
private readonly strategy: IShardingStrategy;
|
||||||
|
|
||||||
public constructor(options: Partial<OptionalWebSocketManagerOptions> & RequiredWebSocketManagerOptions) {
|
public constructor(options: Partial<OptionalWebSocketManagerOptions> & RequiredWebSocketManagerOptions) {
|
||||||
super();
|
super();
|
||||||
this.options = { ...DefaultWebSocketManagerOptions, ...options };
|
this.options = { ...DefaultWebSocketManagerOptions, ...options };
|
||||||
}
|
this.strategy = this.options.buildStrategy(this);
|
||||||
|
|
||||||
public setStrategy(strategy: IShardingStrategy) {
|
|
||||||
this.strategy = strategy;
|
|
||||||
return this;
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
|
|
Loading…
Reference in a new issue