134 lines
4.0 KiB
TypeScript
134 lines
4.0 KiB
TypeScript
import { Collection } from '../utils/collection.ts'
|
|
import type { Client } from './client.ts'
|
|
import { RESTManager } from '../rest/mod.ts'
|
|
import { Gateway } from '../gateway/mod.ts'
|
|
import { HarmonyEventEmitter } from '../utils/events.ts'
|
|
import { GatewayEvents } from '../types/gateway.ts'
|
|
import { delay } from '../utils/delay.ts'
|
|
|
|
// eslint-disable-next-line @typescript-eslint/consistent-type-definitions
|
|
export type ShardManagerEvents = {
|
|
launch: [number]
|
|
shardReady: [number]
|
|
shardDisconnect: [number, number | undefined, string | undefined]
|
|
shardError: [number, Error, ErrorEvent]
|
|
shardResume: [number]
|
|
}
|
|
|
|
export class ShardManager extends HarmonyEventEmitter<ShardManagerEvents> {
|
|
list: Collection<string, Gateway> = new Collection()
|
|
client: Client
|
|
cachedShardCount?: number
|
|
queueProcessing: boolean = false
|
|
queue: CallableFunction[] = []
|
|
|
|
get rest(): RESTManager {
|
|
return this.client.rest
|
|
}
|
|
|
|
constructor(client: Client) {
|
|
super()
|
|
this.client = client
|
|
}
|
|
|
|
debug(msg: string): void {
|
|
this.client.debug('Shards', msg)
|
|
}
|
|
|
|
enqueueIdentify(fn: CallableFunction): ShardManager {
|
|
this.queue.push(fn)
|
|
// eslint-disable-next-line @typescript-eslint/no-floating-promises
|
|
if (!this.queueProcessing) this.processQueue()
|
|
return this
|
|
}
|
|
|
|
private async processQueue(): Promise<void> {
|
|
if (this.queueProcessing || this.queue.length === 0) return
|
|
this.queueProcessing = true
|
|
const item = this.queue[0]
|
|
await item()
|
|
this.queue.shift()
|
|
await delay(5000)
|
|
this.queueProcessing = false
|
|
if (this.queue.length === 0) {
|
|
this.queueProcessing = false
|
|
} else {
|
|
await this.processQueue()
|
|
}
|
|
}
|
|
|
|
async getShardCount(): Promise<number> {
|
|
let shardCount: number
|
|
if (this.cachedShardCount !== undefined) shardCount = this.cachedShardCount
|
|
else {
|
|
if (
|
|
this.client.shardCount === 'auto' &&
|
|
this.client.fetchGatewayInfo !== false
|
|
) {
|
|
this.debug('Fetch /gateway/bot...')
|
|
const info = await this.client.rest.api.gateway.bot.get()
|
|
this.debug(`Recommended Shards: ${info.shards}`)
|
|
this.debug('=== Session Limit Info ===')
|
|
this.debug(
|
|
`Remaining: ${info.session_start_limit.remaining}/${info.session_start_limit.total}`
|
|
)
|
|
this.debug(`Reset After: ${info.session_start_limit.reset_after}ms`)
|
|
shardCount = info.shards as number
|
|
} else
|
|
shardCount =
|
|
typeof this.client.shardCount === 'string'
|
|
? 1
|
|
: this.client.shardCount ?? 1
|
|
}
|
|
this.cachedShardCount = shardCount
|
|
return this.cachedShardCount
|
|
}
|
|
|
|
/** Launches a new Shard */
|
|
async launch(id: number): Promise<ShardManager> {
|
|
if (this.list.has(id.toString()) === true)
|
|
throw new Error(`Shard ${id} already launched`)
|
|
|
|
this.debug(`Launching Shard: ${id}`)
|
|
const shardCount = await this.getShardCount()
|
|
|
|
const gw = new Gateway(this.client, [Number(id), shardCount])
|
|
this.list.set(id.toString(), gw)
|
|
|
|
gw.initWebsocket()
|
|
this.emit('launch', id)
|
|
|
|
gw.on(GatewayEvents.Ready, () => this.emit('shardReady', id))
|
|
gw.on('error', (err: Error, evt: ErrorEvent) =>
|
|
this.emit('shardError', id, err, evt)
|
|
)
|
|
gw.on(GatewayEvents.Resumed, () => this.emit('shardResume', id))
|
|
gw.on('close', (code: number, reason: string) =>
|
|
this.emit('shardDisconnect', id, code, reason)
|
|
)
|
|
|
|
return gw.waitFor(GatewayEvents.Ready, () => true).then(() => this)
|
|
}
|
|
|
|
/** Launches all Shards */
|
|
async connect(): Promise<ShardManager> {
|
|
const shardCount = await this.getShardCount()
|
|
this.client.shardCount = shardCount
|
|
this.debug(`Launching ${shardCount} shard${shardCount === 1 ? '' : 's'}...`)
|
|
const startTime = Date.now()
|
|
for (let i = 0; i < shardCount; i++) {
|
|
await this.launch(i)
|
|
}
|
|
const endTime = Date.now()
|
|
const diff = endTime - startTime
|
|
this.debug(
|
|
`Launched ${shardCount} shards! Time taken: ${Math.floor(diff / 1000)}s`
|
|
)
|
|
return this
|
|
}
|
|
|
|
get(id: number): Gateway | undefined {
|
|
return this.list.get(id.toString())
|
|
}
|
|
}
|