feat(gateway): Add transport compression support (#3815)

* Add zlib-stream transport-compression support * remove external buffer, add the url param this also fixes the CI Deno error * use a queue and promises instead of calling handleDiscordPacket from the callback of inflate.on('data') * Add zstd support * Re-add payload compression, set fzstd as optional, refactor * add comment for decompressionPromisesQueue * fix comment for isCompressed * Don't create a new Uint8Array if the input is a Buffer * Use a Uint8Array buffer instead of parsing incomplete JSON chunks * remove old commented code * Change ts expect error message * use type import for ShardGatewayConfig * Apply suggestions from code review Co-authored-by: LTS (Link) <lts20050703@gmail.com> --------- Co-authored-by: LTS (Link) <lts20050703@gmail.com>
2026-09-17 08:47:22 +00:00 · 2024-10-05 19:48:05 +02:00
parent 57fa3d4d22
commit 7303a20d7a
6 changed files with 227 additions and 13 deletions
@@ -33,6 +33,9 @@
    "@discordeno/utils": "19.0.0-beta.1",
    "ws": "^8.18.0"
  },
+  "optionalDependencies": {
+    "fzstd": "^0.1.1"
+  },
  "devDependencies": {
    "@biomejs/biome": "^1.8.3",
    "@swc/cli": "^0.4.0",
@@ -1,7 +1,9 @@
-import { inflateSync } from 'node:zlib'
+import { Buffer } from 'node:buffer'
+import { Inflate, createInflate, inflateSync, constants as zlibConstants } from 'node:zlib'
 import type { DiscordGatewayPayload, DiscordHello, DiscordReady } from '@discordeno/types'
 import { GatewayCloseEventCodes, GatewayOpcodes } from '@discordeno/types'
 import { LeakyBucket, camelize, delay, logger } from '@discordeno/utils'
+import type { Decompress as ZstdDecompress } from 'fzstd'
 import NodeWebSocket from 'ws'
 import {
  type BotStatusUpdate,
@@ -11,8 +13,18 @@ import {
  ShardSocketCloseCodes,
  type ShardSocketRequest,
  ShardState,
+  TransportCompression,
 } from './types.js'

+const ZLIB_SYNC_FLUSH = new Uint8Array([0x0, 0x0, 0xff, 0xff])
+
+let fzstd: typeof import('fzstd')
+
+/** Since fzstd is an optional dependency, we need to import it lazily */
+async function getFZStd() {
+  return (fzstd ??= await import('fzstd'))
+}
+
 export class DiscordenoShard {
  /** The id of the shard */
  id: number
@@ -44,6 +56,16 @@ export class DiscordenoShard {
  bucket: LeakyBucket
  /** Logger for the bucket */
  logger: Pick<typeof logger, 'debug' | 'info' | 'warn' | 'error' | 'fatal'>
+  /** Text decoder used for compressed payloads */
+  textDecoder = new TextDecoder()
+  /** ZLib Inflate instance for ZLib-stream transport payloads */
+  inflate?: Inflate
+  /** ZLib inflate buffer */
+  inflateBuffer: Uint8Array | null = null
+  /** ZStd Decompress instance for ZStd-stream transport payloads */
+  zstdDecompress?: ZstdDecompress
+  /** Queue for compressed payloads for Zstd Decompress */
+  decompressionPromisesQueue: ((data: DiscordGatewayPayload) => void)[] = []

  constructor(options: ShardCreateOptions) {
    this.id = options.id
@@ -118,6 +140,60 @@ export class DiscordenoShard {
    url.searchParams.set('v', this.gatewayConfig.version.toString())
    url.searchParams.set('encoding', 'json')

+    // Set the compress url param and initialize the decompression contexts
+    if (this.gatewayConfig.transportCompression) {
+      url.searchParams.set('compress', this.gatewayConfig.transportCompression)
+
+      if (this.gatewayConfig.transportCompression === TransportCompression.zlib) {
+        this.inflateBuffer = null
+        this.inflate = createInflate({
+          finishFlush: zlibConstants.Z_SYNC_FLUSH,
+          chunkSize: 64 * 1024,
+        })
+
+        this.inflate.on('error', (e) => {
+          this.logger.error('The was an error in decompressing a ZLib 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
+        })
+      }
+
+      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')
+
+          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)
+          })
+        }
+      }
+    }
+
+    if (this.gatewayConfig.compress && this.gatewayConfig.transportCompression) {
+      this.logger.warn('[Shard] Payload compression has been disabled since transport compression is enabled as well.')
+      this.gatewayConfig.compress = false
+    }
+
    // We check for built-in WebSocket implementations in Bun or Deno, NodeJS v22 has an implementation too but it seems to be less optimized so for now it is better to use the ws npm package
    const shouldUseBuiltin = Reflect.has(globalThis, 'WebSocket') && (Reflect.has(globalThis, 'Bun') || Reflect.has(globalThis, 'Deno'))

@@ -125,9 +201,12 @@ export class DiscordenoShard {
    const socket: WebSocket = shouldUseBuiltin ? new WebSocket(url) : new NodeWebSocket(url)
    this.socket = socket

+    // By default WebSocket will give us a Blob, this changes it so that it gives us an ArrayBuffer
+    socket.binaryType = 'arraybuffer'
+
    socket.onerror = (event) => this.handleError(event)
-    socket.onclose = async (closeEvent) => await this.handleClose(closeEvent)
-    socket.onmessage = async (messageEvent) => await this.handleMessage(messageEvent)
+    socket.onclose = (closeEvent) => this.handleClose(closeEvent)
+    socket.onmessage = (messageEvent) => this.handleMessage(messageEvent)

    return await new Promise((resolve) => {
      socket.onopen = () => {
@@ -278,7 +357,16 @@ export class DiscordenoShard {
  /** Handle a gateway connection close. */
  async handleClose(close: CloseEvent): Promise<void> {
    this.stopHeartbeating()
-    this.logger.debug(`[Shard] Gateway connection closed with code ${close.code} (${close.reason || '<No reason provided>'}).`)
+
+    // Clear the zlib/zstd data
+    this.inflate = undefined
+    this.zstdDecompress = undefined
+    this.inflateBuffer = null
+    this.decompressionPromisesQueue = []
+
+    this.logger.debug(
+      `[Shard] Gateway connection closed with code ${close.code} (${close.reason || '<No reason provided>'}).`,
+    )

    switch (close.code) {
      case ShardSocketCloseCodes.TestingFinished: {
@@ -340,17 +428,74 @@ export class DiscordenoShard {

  /** Handle an incoming gateway message. */
  async handleMessage(message: MessageEvent): Promise<void> {
-    let preProcessMessage = message.data
+    // The ws npm package will use a Buffer, while the global built-in will use ArrayBuffer
+    const isCompressed = message.data instanceof ArrayBuffer || message.data instanceof Buffer

-    // If message compression is enabled, Discord might send zlib compressed payloads.
-    if (this.gatewayConfig.compress && preProcessMessage instanceof Blob) {
-      preProcessMessage = inflateSync(await preProcessMessage.arrayBuffer()).toString()
+    const data = isCompressed ? await this.decompressPacket(message.data) : (JSON.parse(message.data) as DiscordGatewayPayload)
+
+    // Check if the decompression was not successful
+    if (!data) return
+
+    await this.handleDiscordPacket(data)
+  }
+
+  /**
+   * Decompress a zlib/zstd compressed packet
+   *
+   * @private
+   */
+  async decompressPacket(data: ArrayBuffer | Buffer): Promise<DiscordGatewayPayload | null> {
+    // A buffer is a Uint8Array under the hood. An ArrayBuffer is generic, so we need to create the Uint8Array that uses the whole ArrayBuffer
+    const compressedData: Uint8Array = data instanceof Buffer ? data : new Uint8Array(data)
+
+    if (this.gatewayConfig.transportCompression === TransportCompression.zlib) {
+      if (!this.inflate) {
+        this.logger.fatal('[Shard] zlib-stream transport compression was enabled but no instance of Inflate was found.')
+        return null
+      }
+
+      // Alias, used to avoid some null checks in the Promise constructor
+      const inflate = this.inflate
+
+      const writePromise = new Promise<void>((resolve, reject) => {
+        inflate.write(compressedData, 'binary', (error) => (error ? reject(error) : resolve()))
+      })
+
+      if (!endsWithMarker(compressedData, ZLIB_SYNC_FLUSH)) return null
+
+      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)
    }

-    // Safeguard incase decompression failed to make a string.
-    if (typeof preProcessMessage !== 'string') return
+    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
+      }

-    return await this.handleDiscordPacket(JSON.parse(preProcessMessage) as DiscordGatewayPayload)
+      this.zstdDecompress.push(compressedData)
+
+      const decompressionPromise = new Promise<DiscordGatewayPayload>((r) => this.decompressionPromisesQueue.push(r))
+      return await decompressionPromise
+    }
+
+    if (this.gatewayConfig.compress) {
+      const decompressed = inflateSync(compressedData)
+      const decodedData = this.textDecoder.decode(decompressed)
+
+      return JSON.parse(decodedData)
+    }
+
+    return null
  }

  /** Handles a incoming gateway packet. */
@@ -497,6 +642,7 @@ export class DiscordenoShard {
    // Now the event can be safely forwarded.
    this.events.message?.(this, camelize(packet))
  }
+
  /**
   * Override in order to make the shards presence.
   * async in case devs create the presence based on eg. database values.
@@ -599,6 +745,17 @@ export class DiscordenoShard {
  }
 }

+/** Check if the buffer ends with the marker */
+function endsWithMarker(buffer: Uint8Array, marker: Uint8Array) {
+  if (buffer.length < marker.length) return false
+
+  for (let i = 0; i < marker.length; i++) {
+    if (buffer[buffer.length - marker.length + i] !== marker[i]) return false
+  }
+
+  return true
+}
+
 export interface ShardCreateOptions {
  /** The shard id */
  id: number
@@ -17,6 +17,7 @@ import {
  ShardSocketCloseCodes,
  type ShardSocketRequest,
  type StatusUpdate,
+  type TransportCompression,
  type UpdateVoiceState,
 } from './types.js'

@@ -35,6 +36,7 @@ export function createGatewayManager(options: CreateGatewayManagerOptions): Gate
  const gateway: GatewayManager = {
    events: options.events ?? {},
    compress: options.compress ?? false,
+    transportCompression: options.transportCompression ?? null,
    intents: options.intents ?? 0,
    properties: {
      os: options.properties?.os ?? process.platform,
@@ -148,6 +150,7 @@ export function createGatewayManager(options: CreateGatewayManagerOptions): Gate
          id: shardId,
          connection: {
            compress: gateway.compress,
+            transportCompression: gateway.transportCompression ?? null,
            intents: gateway.intents,
            properties: gateway.properties,
            token: gateway.token,
@@ -380,6 +383,7 @@ export function createGatewayManager(options: CreateGatewayManagerOptions): Gate
          id: shardId,
          connection: {
            compress: this.compress,
+            transportCompression: gateway.transportCompression,
            intents: this.intents,
            properties: this.properties,
            token: this.token,
@@ -606,6 +610,8 @@ export interface CreateGatewayManagerOptions {
   * @default false
   */
  compress?: boolean
+  /** What transport compression should be used */
+  transportCompression?: TransportCompression | null
  /** The calculated intent value of the events which the shard should receive.
   *
   * @default 0
@@ -18,12 +18,47 @@ export enum ShardState {
  Offline = 6,
 }

+export enum TransportCompression {
+  /**
+   * ZLib-Stream Transport Compression.
+   *
+   * @remarks
+   * Uses `node:zlib` to decompress the payloads
+   *
+   * @see https://discord.com/developers/docs/topics/gateway#zlibstream
+   */
+  zlib = 'zlib-stream',
+  /**
+   * 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.
+   *
+   * @see https://discord.com/developers/docs/topics/gateway#zstdstream
+   */
+  zstd = 'zstd-stream',
+}
+
 export interface ShardGatewayConfig {
-  /** Whether incoming payloads are compressed using zlib.
+  /**
+   * Whatever to enable Payload compression.
+   *
+   * @remarks
+   * This is compatible with {@link transportCompression}
   *
   * @default false
+   *
+   * @see https://discord.com/developers/docs/topics/gateway#payload-compression
   */
  compress: boolean
+  /**
+   * What Transport Compression should be use
+   *
+   * @default null
+   *
+   * @see https://discord.com/developers/docs/topics/gateway#transport-compression
+   */
+  transportCompression: TransportCompression | null
  /** The calculated intent value of the events which the shard should receive.
   *
   * @default 0
@@ -120,7 +155,8 @@ export enum ShardSocketCloseCodes {
  Shutdown = 3000,
  /** A resume has been requested and therefore the old connection needs to be closed. */
  ResumeClosingOldConnection = 3024,
-  /** Did not receive a heartbeat ACK in time.
+  /**
+   * Did not receive a heartbeat ACK in time.
   * Closing the shard and creating a new session.
   */
  ZombiedConnection = 3010,