Support zstd-stream gateway compression

guides/advanced/gateway_compression.md

@@ -0,0 +1,50 @@
+# Gateway Compression
+
+Nostrum supports either the `zlib-stream` or `zstd-stream` gateway compression
+methods, as documented
+[here](https://discord.com/developers/docs/topics/gateway#encoding-and-compression)
+
+Most users are fine to leave the `gateway_compression` configuration option set
+to `:zlib` (default), but users looking for a potential reduction in payload
+sizes from the Discord gateway can optionally set `:zstd` here.
+
+## Using `:zstd` compression
+
+Using `:zstd` depends on the [`:ezstd`](https://hex.pm/packages/ezstd) library,
+so you will have to add this dependency to your `mix.exs` file:
+
+```elixir
+ defp deps do
+ [
+ {:nostrum, ...},
+ {:ezstd, "~> 1.1"} # new dependency
+ ]
+ end
+```
+
+
+> #### `:ezstd` NIFs {: .info}
+>
+> Some functionality of `:ezstd` depends on Erlang NIFs (Natively Implemented
+> Functions). This means that a proper compiler installation as well as other
+> build tools like `git` may be necessary at the stage where you compile your
+> dependencies.
+>
+> It may be useful to run `mix deps.compile` in any build systems to ensure that
+> your application does not need build utilities in the built application image.
+
+Once you have this additional dependency installed in your project, set the
+`:nostrum`, `:gateway_compression` configuration option to `:zstd` and Nostrum
+should pick up on it.
+
+You will need to run `mix deps.get` and `mix deps.compile` to install and
+compile the new `:ezstd` dependency.
+
+> #### Nostrum detection of `:ezstd` {: .tip}
+>
+> Since the check for `:ezstd` takes place when you compile Nostrum, you might
+> need to run `mix deps.compile --force nostrum` to ensure that Nostrum is
+> recompiled and recognises the newly installed `:ezstd` dependency.
+>
+> Not doing this may mean that your compiled Nostrum version is still using
+> dummy handlers that will error out even when `:ezstd` is installed.

guides/intro/intro.md

@@ -52,6 +52,9 @@ Apart from the `token` field mentioned above, the following fields are also supp
  of members for all guilds at startup. Depending on your [cache
  backend](../advanced/pluggable_caching.md), this may increase startup time
  and memory usage by quite a bit. Defaults to `false`.
+- `gateway_compression` - use either `:zlib` (default) or `:zstd` for compression
+ of messages from the Discord gateway. See the documentation on
+ [Gateway Compression](../advanced/gateway_compression.md) for more information.
 
 
 ### Voice-specific

lib/nostrum/shard/session.ex

@@ -46,12 +46,33 @@ defmodule Nostrum.Shard.Session do
  alias Nostrum.Shard.{Connector, Event, Payload}
  alias Nostrum.Struct.WSState
 
+ alias Nostrum.Shard.Session.Compression
+
  require Logger
 
  @behaviour :gen_statem
 
  # Query string to connect to when upgrading the connection.
- @gateway_qs "/?compress=zlib-stream&encoding=etf&v=10"
+ @gateway_qs "/?encoding=etf&v=10"
+
+ @gateway_compress Application.compile_env(
+ :nostrum,
+ :gateway_compression,
+ :zlib
+ )
+
+ @compression_module (case @gateway_compress do
+ :zlib ->
+ Compression.Zlib
+
+ :zstd ->
+ Compression.Zstd.check_available!()
+ Compression.Zstd
+
+ _ ->
+ raise ArgumentError,
+ "Unsupported compression type: #{@gateway_compress}"
+ end)
 
  # Maximum time the initial connection may take.
  @timeout_connect :timer.seconds(5)
@@ -190,34 +211,44 @@ defmodule Nostrum.Shard.Session do
  # end
 
  def connecting_ws(:enter, _from, %{conn: conn} = data) do
- Logger.debug("Upgrading connection to websocket")
+ Logger.debug("Upgrading connection to websocket with #{@gateway_compress} compression")
  set_timeout = {:state_timeout, @timeout_ws_upgrade, :upgrade_timeout}
- stream = :gun.ws_upgrade(conn, @gateway_qs, [], %{flow: @standard_flow})
+ stream = :gun.ws_upgrade(conn, @gateway_qs <> compression_qs(), [], %{flow: @standard_flow})
  {:keep_state, %{data | stream: stream}, set_timeout}
  end
 
  def connecting_ws(
  :info,
  {:gun_upgrade, _conn, _stream, ["websocket"], _headers},
- %{zlib_ctx: nil} = data
+ %{compress_ctx: nil} = data
  ) do
- zlib_context = :zlib.open()
- :zlib.inflateInit(zlib_context)
+ context = @compression_module.create_context()
 
  {:next_state, :connected,
- %{data | zlib_ctx: zlib_context, last_heartbeat_ack: DateTime.utc_now(), heartbeat_ack: true}}
+ %{
+ data
+ | compress_ctx: context,
+ last_heartbeat_ack: DateTime.utc_now(),
+ heartbeat_ack: true
+ }}
  end
 
  def connecting_ws(
  :info,
  {:gun_upgrade, _conn, _stream, ["websocket"], _headers},
- %{zlib_ctx: zlib_ctx} = data
+ %{compress_ctx: compress_ctx} = data
  ) do
  Logger.info("Re-established websocket connection")
- :ok = :zlib.inflateReset(zlib_ctx)
+
+ compress_ctx = @compression_module.reset_context(compress_ctx)
 
  {:next_state, :connected,
- %{data | last_heartbeat_ack: DateTime.utc_now(), heartbeat_ack: true}}
+ %{
+ data
+ | last_heartbeat_ack: DateTime.utc_now(),
+ heartbeat_ack: true,
+ compress_ctx: compress_ctx
+ }}
  end
 
  def connecting_ws(:state_timeout, :upgrade_timeout, _data) do
@@ -248,9 +279,10 @@ defmodule Nostrum.Shard.Session do
  end
 
  def connected(:info, {:gun_ws, _worker, stream, {:binary, frame}}, data) do
+ payload_decompressed = @compression_module.inflate(data.compress_ctx, frame)
+
  payload =
- data.zlib_ctx
- |> :zlib.inflate(frame)
+ payload_decompressed
  |> :erlang.iolist_to_binary()
  |> :erlang.binary_to_term()
 
@@ -375,4 +407,8 @@ defmodule Nostrum.Shard.Session do
  :timeout
  end
  end
+
+ defp compression_qs do
+ "&compress=" <> Atom.to_string(@gateway_compress) <> "-stream"
+ end
 end

lib/nostrum/shard/session/compression.ex

@@ -0,0 +1,32 @@
+defmodule Nostrum.Shard.Session.Compression do
+ @moduledoc """
+ A behaviour for compression methods supported by the Discord gateway to implement.
+
+ See the modules nested under this behaviour for reference implementations.
+
+ > ### Internal module {: .info}
+ >
+ > This module is intended for exclusive usage inside of nostrum, and is
+ > documented for completeness and people curious to look behind the covers.
+ """
+
+ @doc """
+ Create a new compression context that can be passed as an argument to other
+ methods within the behaviour to inflate data or reset the context to a
+ blank state.
+ """
+ @callback create_context() :: reference()
+
+ @doc """
+ Decompress a frame received from Discord over the gateway. Should return an
+ iolist of the decompressed data.
+ """
+ @callback inflate(reference(), iodata()) :: iolist()
+
+ @doc """
+ Reset a decompression context to a blank slate, this is useful after a websocket
+ resume has taken place or something similar requiring the reset of the state
+ for a shard.
+ """
+ @callback reset_context(reference()) :: reference()
+end

lib/nostrum/shard/session/compression/zlib.ex

@@ -0,0 +1,32 @@
+defmodule Nostrum.Shard.Session.Compression.Zlib do
+ @moduledoc """
+ Implementation of compression methods for the `zlib` compression algorithm.
+
+ > ### Internal module {: .info}
+ >
+ > This module is intended for exclusive usage inside of nostrum, and is
+ > documented for completeness and people curious to look behind the covers.
+ """
+
+ @behaviour Nostrum.Shard.Session.Compression
+
+ @spec create_context() :: :zlib.zstream()
+ def create_context do
+ context = :zlib.open()
+ :zlib.inflateInit(context)
+
+ context
+ end
+
+ @spec inflate(:zlib.zstream(), iodata()) :: iolist()
+ def inflate(ctx, frame) do
+ :zlib.inflate(ctx, frame)
+ end
+
+ @spec reset_context(:zlib.zstream()) :: :zlib.zstream()
+ def reset_context(ctx) do
+ :zlib.inflateReset(ctx)
+
+ ctx
+ end
+end

lib/nostrum/shard/session/compression/zstd.ex

@@ -0,0 +1,61 @@
+defmodule Nostrum.Shard.Session.Compression.Zstd do
+ @moduledoc """
+ Implementation of compression methods for the `zstd` compression algorithm.
+
+ > ### Internal module {: .info}
+ >
+ > This module is intended for exclusive usage inside of nostrum, and is
+ > documented for completeness and people curious to look behind the covers.
+ """
+
+ @behaviour Nostrum.Shard.Session.Compression
+
+ def check_available! do
+ if not Code.ensure_loaded?(:ezstd) do
+ zstd_missing()
+ end
+ end
+
+ @spec zstd_missing() :: no_return
+ defp zstd_missing do
+ raise ArgumentError, """
+ Cannot use the :zstd gateway compression option without optional dependency :ezstd.
+
+ See https://kraigie.github.io/nostrum/gateway_compression.html for more information.
+ """
+ end
+
+ if Code.ensure_loaded?(:ezstd) do
+ @zstd_buffer_size 16_000
+
+ @spec create_context() :: reference() | {:error, any()}
+ def create_context do
+ :ezstd.create_decompression_context(@zstd_buffer_size)
+ end
+
+ @spec inflate(reference(), iodata()) :: iolist() | {:error, any()}
+ def inflate(ctx, frame) do
+ :ezstd.decompress_streaming(ctx, frame)
+ end
+
+ @spec reset_context(reference()) :: reference()
+ def reset_context(_ctx) do
+ create_context()
+ end
+ else
+ @spec create_context() :: reference() | {:error, any()}
+ def create_context do
+ zstd_missing()
+ end
+
+ @spec inflate(reference(), iodata()) :: iodata() | {:error, any()}
+ def inflate(_ctx, _frame) do
+ zstd_missing()
+ end
+
+ @spec reset_context(reference()) :: reference()
+ def reset_context(_ctx) do
+ zstd_missing()
+ end
+ end
+end

lib/nostrum/struct/ws_state.ex

@@ -17,7 +17,7 @@ defmodule Nostrum.Struct.WSState do
  :last_heartbeat_ack,
  :heartbeat_ack,
  :heartbeat_interval,
- :zlib_ctx
+ :compress_ctx
  ]
 
  @typedoc "The shard number"
@@ -73,8 +73,8 @@ defmodule Nostrum.Struct.WSState do
  @typedoc "Interval at which heartbeats are sent"
  @type heartbeat_interval :: pos_integer() | nil
 
- @typedoc "Reference to the current zlib context"
- @type zlib_ctx :: reference | nil
+ @typedoc "Reference to the current compression context"
+ @type compress_ctx :: reference | nil
 
  @type t :: %__MODULE__{
  shard_num: shard_num,
@@ -90,6 +90,6 @@ defmodule Nostrum.Struct.WSState do
  last_heartbeat_ack: last_heartbeat_ack,
  heartbeat_ack: heartbeat_ack,
  heartbeat_interval: heartbeat_interval,
- zlib_ctx: zlib_ctx
+ compress_ctx: compress_ctx
  }
 end

mix.exs

@@ -81,6 +81,7 @@ defmodule Nostrum.Mixfile do
  "guides/advanced/pluggable_caching.md",
  "guides/advanced/multi_node.md",
  "guides/advanced/hot_code_upgrade.md",
+ "guides/advanced/gateway_compression.md",
  "guides/cheat-sheets/api.cheatmd",
  "guides/cheat-sheets/qlc.cheatmd",
  "guides/cheat-sheets/voice.cheatmd"
@@ -159,6 +160,7 @@ defmodule Nostrum.Mixfile do
  {:certifi, "~> 2.13"},
  {:kcl, "~> 1.4"},
  {:mime, "~> 1.6 or ~> 2.0"},
+ {:ezstd, "~> 1.1", optional: true},
  {:castle, "~> 0.3.0", runtime: false},
  {:ex_doc, "~> 0.32", only: :dev, runtime: false},
  {:credo, "~> 1.7.5", only: [:dev, :test], runtime: false},

mix.lock

@@ -15,6 +15,7 @@
  "equivalex": {:hex, :equivalex, "1.0.3", "170d9a82ae066e0020dfe1cf7811381669565922eb3359f6c91d7e9a1124ff74", [:mix], [], "hexpm", "46fa311adb855117d36e461b9c0ad2598f72110ad17ad73d7533c78020e045fc"},
  "erlex": {:hex, :erlex, "0.2.6", "c7987d15e899c7a2f34f5420d2a2ea0d659682c06ac607572df55a43753aa12e", [:mix], [], "hexpm", "2ed2e25711feb44d52b17d2780eabf998452f6efda104877a3881c2f8c0c0c75"},
  "ex_doc": {:hex, :ex_doc, "0.32.1", "21e40f939515373bcdc9cffe65f3b3543f05015ac6c3d01d991874129d173420", [:mix], [{:earmark_parser, "~> 1.4.39", [hex: :earmark_parser, repo: "hexpm", optional: false]}, {:makeup_c, ">= 0.1.1", [hex: :makeup_c, repo: "hexpm", optional: true]}, {:makeup_elixir, "~> 0.14", [hex: :makeup_elixir, repo: "hexpm", optional: false]}, {:makeup_erlang, "~> 0.1", [hex: :makeup_erlang, repo: "hexpm", optional: false]}], "hexpm", "5142c9db521f106d61ff33250f779807ed2a88620e472ac95dc7d59c380113da"},
+ "ezstd": {:hex, :ezstd, "1.1.0", "d3b483d6acfadfb65dba4015371e6d54526dbf3d9ef0941b5add8bf5890731f4", [:rebar3], [], "hexpm", "28cfa0ed6cc3922095ad5ba0f23392a1664273358b17184baa909868361184e7"},
  "file_system": {:hex, :file_system, "1.0.0", "b689cc7dcee665f774de94b5a832e578bd7963c8e637ef940cd44327db7de2cd", [:mix], [], "hexpm", "6752092d66aec5a10e662aefeed8ddb9531d79db0bc145bb8c40325ca1d8536d"},
  "forecastle": {:hex, :forecastle, "0.1.1", "89dcfaccbfffe866cbd8a4c41ade55f62f00f1b5d0528bec787b1e6631004b98", [:mix], [], "hexpm", "f6f4d297224a22ac4387d305249aed7b8b02e85b4a03e83225af4536812c4079"},
  "gun": {:hex, :gun, "2.0.1", "160a9a5394800fcba41bc7e6d421295cf9a7894c2252c0678244948e3336ad73", [:make, :rebar3], [{:cowlib, "2.12.1", [hex: :cowlib, repo: "hexpm", optional: false]}], "hexpm", "a10bc8d6096b9502205022334f719cc9a08d9adcfbfc0dbee9ef31b56274a20b"},