From d79e672f5a5818f755a3654066741b21e8dbc286 Mon Sep 17 00:00:00 2001 From: Fleny Date: Sun, 17 Aug 2025 22:38:46 +0200 Subject: [PATCH] feat(gateway): Use `node:zlib` zstd when available (#4373) Co-authored-by: Link --- packages/gateway/src/Shard.ts | 101 +++++++++++++++++++++++++--------- packages/gateway/src/types.ts | 4 +- 2 files changed, 79 insertions(+), 26 deletions(-) diff --git a/packages/gateway/src/Shard.ts b/packages/gateway/src/Shard.ts index e9a0ef667..6d6ff6523 100644 --- a/packages/gateway/src/Shard.ts +++ b/packages/gateway/src/Shard.ts @@ -1,9 +1,9 @@ import { Buffer } from 'node:buffer' -import { createInflate, type Inflate, inflateSync, constants as zlibConstants } from 'node:zlib' +import zlib from 'node:zlib' import type { DiscordGatewayPayload, DiscordHello, DiscordReady, DiscordUpdatePresence } from '@discordeno/types' import { GatewayCloseEventCodes, GatewayOpcodes } from '@discordeno/types' import { camelize, delay, LeakyBucket, logger } from '@discordeno/utils' -import type { Decompress as ZstdDecompress } from 'fzstd' +import type { Decompress as FZstdDecompress } from 'fzstd' import NodeWebSocket from 'ws' import { type ShardEvents, @@ -67,12 +67,12 @@ export class DiscordenoShard { goingOffline = false /** Text decoder used for compressed payloads. */ textDecoder = new TextDecoder() - /** ZLib Inflate instance for ZLib-stream transport payloads. */ - inflate?: Inflate + /** zlib Inflate or zstd decompress (from node:zlib) instance for transport payloads. */ + inflate?: zlib.Inflate | zlib.ZstdDecompress /** ZLib inflate buffer. */ inflateBuffer: Uint8Array | null = null /** ZStd Decompress instance for ZStd-stream transport payloads. */ - zstdDecompress?: ZstdDecompress + zstdDecompress?: FZstdDecompress /** Queue for compressed payloads for Zstd Decompress */ decompressionPromisesQueue: ((data: DiscordGatewayPayload) => void)[] = [] /** @@ -185,8 +185,8 @@ export class DiscordenoShard { if (this.gatewayConfig.transportCompression === TransportCompression.zlib) { this.inflateBuffer = null - this.inflate = createInflate({ - finishFlush: zlibConstants.Z_SYNC_FLUSH, + this.inflate = zlib.createInflate({ + finishFlush: zlib.constants.Z_SYNC_FLUSH, chunkSize: 64 * 1024, }) @@ -211,19 +211,49 @@ export class DiscordenoShard { } if (this.gatewayConfig.transportCompression === TransportCompression.zstd) { - const fzstd = await getFZStd().catch(() => { - this.logger.warn('[Shard] "fzstd" is not installed. Disabled transport compression.') - url.searchParams.delete('compress') + if ('createZstdDecompress' in zlib) { + this.logger.debug('[Shard] Using node:zlib zstd decompression.') - return null - }) - - if (fzstd) { - this.zstdDecompress = new fzstd.Decompress((data) => { - const decodedData = this.textDecoder.decode(data) - const parsedData = JSON.parse(decodedData) - this.decompressionPromisesQueue.shift()?.(parsedData) + this.inflateBuffer = null + this.inflate = zlib.createZstdDecompress({ + chunkSize: 64 * 1024, }) + + this.inflate.on('error', (e) => { + this.logger.error('The was an error in decompressing a Zstd compressed payload', e) + }) + + this.inflate.on('data', (data) => { + if (!(data instanceof Uint8Array)) return + + if (this.inflateBuffer) { + const newBuffer = new Uint8Array(this.inflateBuffer.byteLength + data.byteLength) + newBuffer.set(this.inflateBuffer) + newBuffer.set(data, this.inflateBuffer.byteLength) + this.inflateBuffer = newBuffer + + return + } + + this.inflateBuffer = data + }) + } else { + const fzstd = await getFZStd().catch(() => { + this.logger.warn('[Shard] "fzstd" is not installed. Disabled transport compression.') + url.searchParams.delete('compress') + + return null + }) + + if (fzstd) { + this.logger.debug('[Shard] Using fzstd zstd decompression.') + + this.zstdDecompress = new fzstd.Decompress((data) => { + const decodedData = this.textDecoder.decode(data) + const parsedData = JSON.parse(decodedData) + this.decompressionPromisesQueue.shift()?.(parsedData) + }) + } } } } @@ -547,19 +577,40 @@ export class DiscordenoShard { } if (this.gatewayConfig.transportCompression === TransportCompression.zstd) { - if (!this.zstdDecompress) { - this.logger.fatal('[Shard] zstd-stream transport compression was enabled but no instance of Decompress was found.') - return null + if (this.zstdDecompress) { + this.zstdDecompress.push(compressedData) + + const decompressionPromise = new Promise((r) => this.decompressionPromisesQueue.push(r)) + return await decompressionPromise } - this.zstdDecompress.push(compressedData) + if (this.inflate) { + // Alias, used to avoid some null checks in the Promise constructor + const decompress = this.inflate - const decompressionPromise = new Promise((r) => this.decompressionPromisesQueue.push(r)) - return await decompressionPromise + const writePromise = new Promise((resolve, reject) => { + decompress.write(compressedData, 'binary', (error) => (error ? reject(error) : resolve())) + }) + + await writePromise + + if (!this.inflateBuffer) { + this.logger.warn('[Shard] The ZLib inflate buffer was cleared at an unexpected moment.') + return null + } + + const decodedData = this.textDecoder.decode(this.inflateBuffer) + this.inflateBuffer = null + + return JSON.parse(decodedData) + } + + this.logger.fatal('[Shard] zstd-stream transport compression was enabled but no zstd decompressor was found.') + return null } if (this.gatewayConfig.compress) { - const decompressed = inflateSync(compressedData) + const decompressed = zlib.inflateSync(compressedData) const decodedData = this.textDecoder.decode(decompressed) return JSON.parse(decodedData) diff --git a/packages/gateway/src/types.ts b/packages/gateway/src/types.ts index 7bde05033..e914f855f 100644 --- a/packages/gateway/src/types.ts +++ b/packages/gateway/src/types.ts @@ -32,7 +32,9 @@ export enum TransportCompression { * ZStd-Stream Transport Compression. * * @remarks - * Uses `fzstd` to decompress the payloads. `fzstd` is an optional dependency, it is required to be installed for this compression. + * This will use `node:zlib` zstd implementation if available (node v22.15+ or node v23.8+) or `fzstd` otherwise. + * + * `fzstd` is an optional dependency. You will need to install it to be able to use it. * * @see https://discord.com/developers/docs/topics/gateway#zstdstream */