Skip to content

feat: zstd streaming support #10758

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 4 commits into
base: main
Choose a base branch
from
Draft

feat: zstd streaming support #10758

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
b49529d
feat(ws): zstd streaming support
vladfrangu Feb 14, 2025
d9ce818
chore: naming of the mode
vladfrangu Feb 14, 2025
51bc477
Merge branch 'main' into feat/zstd-streaming
almeidx May 9, 2025
02175c0
fix: remove `@ts-expect-error`
almeidx May 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions packages/ws/src/utils/constants.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ export enum Encoding {
export enum CompressionMethod {
ZlibNative,
ZlibSync,
ZstdNative,
}

export const DefaultDeviceProperty = `@discordjs/ws [VI]{{inject}}[/VI]` as `@discordjs/ws ${string}`;
Expand All @@ -29,6 +30,7 @@ const getDefaultSessionStore = lazy(() => new Collection<number, SessionInfo | n
export const CompressionParameterMap = {
[CompressionMethod.ZlibNative]: 'zlib-stream',
[CompressionMethod.ZlibSync]: 'zlib-stream',
[CompressionMethod.ZstdNative]: 'zstd-stream',
} as const satisfies Record<CompressionMethod, string>;

/**
Expand Down
42 changes: 36 additions & 6 deletions packages/ws/src/ws/WebSocketShard.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -222,7 +222,7 @@ export class WebSocketShard extends AsyncEventEmitter<WebSocketShardEventsMap> {

this.nativeInflate = inflate;
} else {
console.warn('WebSocketShard: Compression is set to native but node:zlib is not available.');
console.warn('WebSocketShard: Compression is set to native zlib but node:zlib is not available.');
params.delete('compress');
}

Expand All @@ -243,6 +243,34 @@ export class WebSocketShard extends AsyncEventEmitter<WebSocketShardEventsMap> {

break;
}

case CompressionMethod.ZstdNative: {
const zlib = await getNativeZlib();
if (zlib && 'createZstdDecompress' in zlib) {
this.inflateBuffer = [];

const inflate = zlib.createZstdDecompress({
chunkSize: 65_535,
}) as nativeZlib.Inflate;
Comment on lines +252 to +254
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I believe this will need to set flush: zlib.constants.ZSTD_e_flush, as the default zstd behaviour is not to flush the output buffer if the received message doesn't close the current frame.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Last I tested this, it worked fine without that 👀

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Well, each ws message is a full payload compressed with zstd, and I've also tested this on a bot in a guild that was getting many payloads/sec with 0 issues. Coulddd you test the hypothesis please and let us know results 🙏

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks like this applies only to the compression stream, not decompression. Ignore me! 👍


inflate.on('data', (chunk) => {
this.inflateBuffer.push(chunk);
});

inflate.on('error', (error) => {
this.emit(WebSocketShardEvents.Error, error);
});

this.nativeInflate = inflate;
} else {
console.warn(
'WebSocketShard: Compression is set to native zstd but node:zlib is not available or your node version does not support zstd decompression.',
);
params.delete('compress');
}

break;
}
}
}

Expand Down Expand Up @@ -633,12 +661,14 @@ export class WebSocketShard extends AsyncEventEmitter<WebSocketShardEventsMap> {

// Deal with transport compression
if (this.transportCompressionEnabled) {
// Each WS message received is a full gateway message for zstd streaming, but for zlib it's chunked
const flush =
decompressable.length >= 4 &&
decompressable.at(-4) === 0x00 &&
decompressable.at(-3) === 0x00 &&
decompressable.at(-2) === 0xff &&
decompressable.at(-1) === 0xff;
this.strategy.options.compression === CompressionMethod.ZstdNative ||
(decompressable.length >= 4 &&
decompressable.at(-4) === 0x00 &&
decompressable.at(-3) === 0x00 &&
decompressable.at(-2) === 0xff &&
decompressable.at(-1) === 0xff);

if (this.nativeInflate) {
const doneWriting = new Promise<void>((resolve) => {
Expand Down