diff --git a/guides/intro/intro.md b/guides/intro/intro.md index d96e8e310..5b25b3a26 100644 --- a/guides/intro/intro.md +++ b/guides/intro/intro.md @@ -45,6 +45,8 @@ Apart from the `token` field mentioned above, the following fields are also supp should contain the total amount of shards that your bot is expected to have. Useful for splitting a single bot across multiple servers, but see also [the multi-node documentation](../advanced/multi_node.md). + - `:manual`: nostrum does not automatically spawn shards. You should use + `Nostrum.Shard.Supervisor.connect/2` to spawn shards instead. - `gateway_intents` - a list of atoms representing gateway intents for Nostrum to subscribe to from the Discord API. More information can be found in the [gateway intents](./gateway_intents.md) documentation page. diff --git a/lib/nostrum/shard.ex b/lib/nostrum/shard.ex index 7f77ad2c7..fca655188 100644 --- a/lib/nostrum/shard.ex +++ b/lib/nostrum/shard.ex @@ -1,14 +1,33 @@ defmodule Nostrum.Shard do @moduledoc false - use Supervisor + use Supervisor, restart: :transient alias Nostrum.Shard.Session - def start_link([_, shard_num, _total] = opts) do + def start_link({:connect, [_, shard_num, _total]} = opts) do Supervisor.start_link(__MODULE__, opts, name: :"Nostrum.Shard-#{shard_num}") end + def start_link( + {:reconnect, + %{ + shard_num: shard_num, + total_shards: _total_shards, + gateway: _gateway, + resume_gateway: _resume_gateway, + seq: _seq, + session: _session + }} = + opts + ) do + Supervisor.start_link(__MODULE__, opts, name: :"Nostrum.Shard-#{shard_num}") + end + + def start_link([_, _shard_num, _total] = opts) do + start_link({:connect, opts}) + end + def init(opts) do children = [ {Session, opts} @@ -16,6 +35,11 @@ defmodule Nostrum.Shard do # TODO: Add per shard cache ] - Supervisor.init(children, strategy: :one_for_all, max_restarts: 3, max_seconds: 60) + Supervisor.init(children, + strategy: :one_for_all, + max_restarts: 3, + max_seconds: 60, + auto_shutdown: :any_significant + ) end end diff --git a/lib/nostrum/shard/session.ex b/lib/nostrum/shard/session.ex index 6e3de6243..d882ca1c3 100644 --- a/lib/nostrum/shard/session.ex +++ b/lib/nostrum/shard/session.ex @@ -107,17 +107,45 @@ defmodule Nostrum.Shard.Session do :gen_statem.cast(pid, {:request_guild_members, payload}) end + def disconnect(pid) do + :gen_statem.call(pid, {:disconnect, nil}) + end + + def disconnect(pid, timeout) do + :gen_statem.call(pid, {:disconnect, nil}, timeout) + end + def get_ws_state(pid) do :sys.get_state(pid) end # State machine API + def start_link({:connect, [_gateway, _shard_num, _total]} = opts, statem_opts) do + :gen_statem.start_link(__MODULE__, opts, statem_opts) + end + + def start_link( + {:reconnect, + %{ + shard_num: _shard_num, + total_shards: _total_shards, + gateway: _gateway, + resume_gateway: _resume_gateway, + seq: _seq, + session: _session + }} = + opts, + statem_opts + ) do + :gen_statem.start_link(__MODULE__, opts, statem_opts) + end + def start_link([_gateway, _shard_num, _total] = shard_opts, statem_opts) do :gen_statem.start_link(__MODULE__, shard_opts, statem_opts) end - def init([gateway, shard_num, total]) do + def init({:connect, [gateway, shard_num, total]}) do Logger.metadata(shard: shard_num) state = %WSState{ @@ -131,6 +159,37 @@ defmodule Nostrum.Shard.Session do {:ok, :disconnected, state, connect} end + def init( + {:reconnect, + %{ + shard_num: shard_num, + total_shards: total_shards, + gateway: gateway, + resume_gateway: resume_gateway, + seq: seq, + session: session + }} + ) do + Logger.metadata(shard: shard_num) + + state = %WSState{ + conn_pid: self(), + shard_num: shard_num, + total_shards: total_shards, + gateway: gateway, + resume_gateway: resume_gateway, + session: session, + seq: seq + } + + connect = {:next_event, :internal, :connect} + {:ok, :disconnected, state, connect} + end + + def init([gateway, shard_num, total]) do + init({:connect, [gateway, shard_num, total]}) + end + def callback_mode, do: [:state_functions, :state_enter] def child_spec(opts) do @@ -138,7 +197,8 @@ defmodule Nostrum.Shard.Session do id: __MODULE__, start: {__MODULE__, :start_link, [opts, []]}, type: :worker, - restart: :permanent, + restart: :transient, + significant: true, shutdown: 500 } end @@ -334,6 +394,30 @@ defmodule Nostrum.Shard.Session do :keep_state_and_data end + def connected({:call, from}, {:disconnect, nil}, %{ + conn: conn, + shard_num: shard_num, + total_shards: total, + gateway: gateway, + resume_gateway: resume_gateway, + seq: seq, + session: session + }) do + :ok = :gun.close(conn) + :ok = :gun.flush(conn) + + {:stop_and_reply, :normal, + {:reply, from, + %{ + shard_num: shard_num, + total_shards: total, + gateway: gateway, + resume_gateway: resume_gateway, + session: session, + seq: seq + }}} + end + def connected( :state_timeout, :send_heartbeat = request, diff --git a/lib/nostrum/shard/supervisor.ex b/lib/nostrum/shard/supervisor.ex index c3c9d7b7b..1a2afc70e 100644 --- a/lib/nostrum/shard/supervisor.ex +++ b/lib/nostrum/shard/supervisor.ex @@ -34,7 +34,26 @@ defmodule Nostrum.Shard.Supervisor do handle yet, thus growing the message queue and the memory usage. """ - use Supervisor + @typedoc """ + Shard number(`shard_id`). Range is `0..total_shards-1`. + """ + @type shard_num :: non_neg_integer() + @typedoc """ + Total shard count(`num_shards`). + """ + @type total_shards :: pos_integer() + @typedoc """ + Represents gateway resume information. + """ + @type resume_information :: %{ + shard_num: shard_num(), + total_shards: total_shards(), + gateway: String.t(), + resume_gateway: String.t() | nil, + session: String.t(), + seq: pos_integer() + } + use DynamicSupervisor alias Nostrum.Error.CacheError alias Nostrum.Shard @@ -64,21 +83,33 @@ defmodule Nostrum.Shard.Supervisor do end def start_link(_args) do - {url, gateway_shard_count} = Util.gateway() + {_url, gateway_shard_count} = Util.gateway() - value = Application.get_env(:nostrum, :num_shards, :auto) - shard_range = cast_shard_range(gateway_shard_count, value) + on_start = + DynamicSupervisor.start_link( + __MODULE__, + nil, + name: __MODULE__ + ) - Supervisor.start_link( - __MODULE__, - [url, shard_range], - name: __MODULE__ - ) + case Application.get_env(:nostrum, :num_shards, :auto) do + :manual -> + on_start + + value -> + {lowest, highest, total} = cast_shard_range(gateway_shard_count, value) + + shard_range = lowest..highest + + for num <- shard_range, do: connect(num - 1, total) + end + + on_start end def update_status(status, game, stream, type) do __MODULE__ - |> Supervisor.which_children() + |> DynamicSupervisor.which_children() |> Enum.filter(fn {_id, _pid, _type, [modules]} -> modules == Nostrum.Shard end) |> Enum.map(fn {_id, pid, _type, _modules} -> Supervisor.which_children(pid) end) |> List.flatten() @@ -102,18 +133,55 @@ defmodule Nostrum.Shard.Supervisor do end @doc false - def init([url, {lowest, highest, total}]) do - shard_range = lowest..highest - children = for num <- shard_range, do: create_worker(url, num - 1, total) - - Supervisor.init(children, strategy: :one_for_one, max_restarts: 3, max_seconds: 60) + def init(_) do + DynamicSupervisor.init(strategy: :one_for_one, max_restarts: 3, max_seconds: 60) end @doc false - def create_worker(gateway, shard_num, total) do - Supervisor.child_spec( - {Shard, [gateway, shard_num, total]}, - id: shard_num + def create_worker(shard_num, total) do + {gateway, _gateway_shard_count} = Util.gateway() + {Shard, [gateway, shard_num, total]} + end + + @doc """ + Disconnects the shard with the given shard number from the Gateway. + This function returns `resume_information` given to `Nostrum.Shard.Supervisor.reconnect/1`. + """ + @spec disconnect(shard_num()) :: resume_information() + def disconnect(shard_num) do + :"Nostrum.Shard-#{shard_num}" + |> Supervisor.which_children() + |> Enum.find(fn {id, _pid, _type, _modules} -> id == Nostrum.Shard.Session end) + |> elem(1) + |> Session.disconnect() + end + + @doc """ + Spawns a shard with the specified number and connects it to the discord gateway. + """ + @spec connect(shard_num(), total_shards()) :: DynamicSupervisor.on_start_child() + def connect(shard_num, total_shards) do + DynamicSupervisor.start_child(__MODULE__, create_worker(shard_num, total_shards)) + end + + @doc """ + Reconnect to the gateway using the given `resume_information`. + For more information about resume, please visit [the Discord Developer Portal](https://discord.com/developers/docs/topics/gateway#resuming). + """ + @spec reconnect(resume_information()) :: DynamicSupervisor.on_start_child() + def reconnect( + %{ + shard_num: _shard_num, + total_shards: _total_shards, + gateway: _gateway, + resume_gateway: _resume_gateway, + seq: _seq, + session: _session + } = opts + ) do + DynamicSupervisor.start_child( + __MODULE__, + {Shard, {:reconnect, opts}} ) end end diff --git a/lib/nostrum/util.ex b/lib/nostrum/util.ex index 92fd235cb..435c27494 100644 --- a/lib/nostrum/util.ex +++ b/lib/nostrum/util.ex @@ -83,6 +83,7 @@ defmodule Nostrum.Util do case num do {_lowest, _highest, total} -> total + :manual -> 0 nil -> 1 end end