Repository: phoenixframework/phoenix
Pull request #5363 (2023-03-15 20:23:53)

WebSocket draining


Patch:
diff --git a/installer/templates/phx_web/components/core_components.ex b/installer/templates/phx_web/components/core_components.ex
index fab02c0b4..4823704a5 100644
--- a/installer/templates/phx_web/components/core_components.ex
+++ b/installer/templates/phx_web/components/core_components.ex
@@ -254,7 +254,7 @@ defmodule <%= @web_namespace %>.CoreComponents do
   attr :prompt, :string, default: nil, doc: "the prompt for select inputs"
   attr :options, :list, doc: "the options to pass to Phoenix.HTML.Form.options_for_select/2"
   attr :multiple, :boolean, default: false, doc: "the multiple flag for select inputs"
-  attr :rest, :global, include: ~w(autocomplete cols disabled form max maxlength min minlength
+  attr :rest, :global, include: ~w(autocomplete cols disabled form list max maxlength min minlength
                                    pattern placeholder readonly required rows size step)
   slot :inner_block
 
diff --git a/lib/phoenix/channel/server.ex b/lib/phoenix/channel/server.ex
index bbf7c5413..0e985eef2 100644
--- a/lib/phoenix/channel/server.ex
+++ b/lib/phoenix/channel/server.ex
@@ -18,14 +18,14 @@ defmodule Phoenix.Channel.Server do
   def join(socket, channel, message, opts) do
     %{topic: topic, payload: payload, ref: ref, join_ref: join_ref} = message
 
-    starter = opts[:starter] || &PoolSupervisor.start_child/4
+    starter = opts[:starter] || &PoolSupervisor.start_child/3
     assigns = Map.merge(socket.assigns, Keyword.get(opts, :assigns, %{}))
     socket = %{socket | topic: topic, channel: channel, join_ref: join_ref || ref, assigns: assigns}
     ref = make_ref()
     from = {self(), ref}
     child_spec = channel.child_spec({socket.endpoint, from})
 
-    case starter.(socket.endpoint, socket.handler, from, child_spec) do
+    case starter.(socket, from, child_spec) do
       {:ok, pid} ->
         send(pid, {Phoenix.Channel, payload, from, socket})
         mon_ref = Process.monitor(pid)
@@ -321,6 +321,14 @@ defmodule Phoenix.Channel.Server do
     handle_in(result)
   end
 
+  def handle_info(
+        %Broadcast{event: "phx_drain"},
+        %{transport_pid: transport_pid} = socket
+      ) do
+    send(transport_pid, :socket_drain)
+    {:stop, {:shutdown, :draining}, socket}
+  end
+
   def handle_info(
         %Broadcast{topic: topic, event: event, payload: payload},
         %Socket{topic: topic} = socket
diff --git a/lib/phoenix/endpoint.ex b/lib/phoenix/endpoint.ex
index 479df96ec..bf23635a7 100644
--- a/lib/phoenix/endpoint.ex
+++ b/lib/phoenix/endpoint.ex
@@ -232,11 +232,15 @@ defmodule Phoenix.Endpoint do
       as defined by [`Plug.Cowboy`](https://hexdocs.pm/plug_cowboy/). Defaults
       to `false`
 
-    * `:drainer` - a drainer process that triggers when your application is
-      shutting down to wait for any on-going request to finish. It accepts all
-      options as defined by [`Plug.Cowboy.Drainer`](https://hexdocs.pm/plug_cowboy/Plug.Cowboy.Drainer.html).
-      Defaults to `[]`, which will start a drainer process for each configured endpoint,
-      but can be disabled by setting it to `false`.
+    * `:drainer` - a drainer process waits for any on-going request to finish
+      during application shutdown. It accepts the `:shutdown` and
+      `:drain_check_interval` options as defined by `Plug.Cowboy.Drainer`.
+      Note the draining does not terminate any existing connection, it simply
+      waits for them to finish. Socket connections run their own drainer
+      before this one is invoked. That's because sockets are stateful and
+      can be gracefully notified, which allows us to stagger them over a
+      longer period of time. See the documentation for `socket/3` for more
+      information
 
   ## Endpoint API
 
@@ -790,57 +794,75 @@ defmodule Phoenix.Endpoint do
     * `:code_reloader` - enable or disable the code reloader. Defaults to your
       endpoint configuration
 
-    * `:connect_info` - a list of keys that represent data to be copied from
-      the transport to be made available in the user socket `connect/3` callback
+    * `:drainer` - a keyword list configuring how to drain sockets
+      on application shutdown. The goal is to notify all channels (and
+      LiveViews) clients to reconnect. The supported options are:
 
-      The valid keys are:
+      * `:shutdown` - How long to wait for connections to drain. Defaults to 15000ms.
+      * `:drain_check_interval` - The maximum frequency to terminate batches. Defaults to 1000ms.
 
-        * `:peer_data` - the result of `Plug.Conn.get_peer_data/1`
+      For example, if you have 150k connections, the default values will
+      split them into 15 batches of 10k connections and notify each batch
+      within 1000ms. If the batch shutdowns faster, we move on to the next
+      one before the interval. If it takes longer, then we proceed anyway.
+      Note that, after the socket drainer runs, the lower level HTTP/HTTPS
+      connection drainer will still run, and apply to all connections.
+      Set it to `false` to disable draining.
 
-        * `:trace_context_headers` - a list of all trace context headers. Supported
-          headers are defined by the [W3C Trace Context Specification](https://www.w3.org/TR/trace-context-1/).
-          These headers are necessary for libraries such as [OpenTelemetry](https://opentelemetry.io/) to extract
-          trace propagation information to know this request is part of a larger trace
-          in progress.
+    * `:connect_info` - a list of keys that represent data to be copied from
+      the transport to be made available in the user socket `connect/3` callback.
+      See the "Connect info" subsection for valid keys
 
-        * `:x_headers` - all request headers that have an "x-" prefix
+  ### Connect info
 
-        * `:uri` - a `%URI{}` with information from the conn
+  The valid keys are:
 
-        * `:user_agent` - the value of the "user-agent" request header
+    * `:peer_data` - the result of `Plug.Conn.get_peer_data/1`
 
-        * `{:session, session_config}` - the session information from `Plug.Conn`.
-          The `session_config` is an exact copy of the arguments given to `Plug.Session`.
-          This requires the "_csrf_token" to be given as request parameter with
-          the value of `URI.encode_www_form(Plug.CSRFProtection.get_csrf_token())`
-          when connecting to the socket. It can also be a MFA to allow loading
-          config in runtime `{MyAppWeb.Auth, :get_session_config, []}`. Otherwise
-          the session will be `nil`.
+    * `:trace_context_headers` - a list of all trace context headers. Supported
+      headers are defined by the [W3C Trace Context Specification](https://www.w3.org/TR/trace-context-1/).
+      These headers are necessary for libraries such as [OpenTelemetry](https://opentelemetry.io/) to extract
+      trace propagation information to know this request is part of a larger trace
+      in progress.
 
-          `session_config` may take a `:csrf_token_key` option
-          which is useful when using `:protect_from_forgery` with a custom
-          `:session_key`. If not given, it defaults to `"_csrf_token"`.
+    * `:x_headers` - all request headers that have an "x-" prefix
 
-      Arbitrary keywords may also appear following the above valid keys, which
-      is useful for passing custom connection information to the socket.
+    * `:uri` - a `%URI{}` with information from the conn
 
-      For example:
+    * `:user_agent` - the value of the "user-agent" request header
 
-      ```
-        socket "/socket", AppWeb.UserSocket,
-            websocket: [
-              connect_info: [:peer_data, :trace_context_headers, :x_headers, :uri, session: [store: :cookie]]
-            ]
-      ```
+    * `{:session, session_config}` - the session information from `Plug.Conn`.
+      The `session_config` is an exact copy of the arguments given to `Plug.Session`.
+      This requires the "_csrf_token" to be given as request parameter with
+      the value of `URI.encode_www_form(Plug.CSRFProtection.get_csrf_token())`
+      when connecting to the socket. It can also be a MFA to allow loading
+      config in runtime `{MyAppWeb.Auth, :get_session_config, []}`. Otherwise
+      the session will be `nil`.
 
-      With arbitrary keywords:
+      `session_config` may take a `:csrf_token_key` option
+      which is useful when using `:protect_from_forgery` with a custom
+      `:session_key`. If not given, it defaults to `"_csrf_token"`.
 
-      ```
-        socket "/socket", AppWeb.UserSocket,
-            websocket: [
-              connect_info: [:uri, custom_value: "abcdef"]
-            ]
-      ```
+  Arbitrary keywords may also appear following the above valid keys, which
+  is useful for passing custom connection information to the socket.
+
+  For example:
+
+  ```
+    socket "/socket", AppWeb.UserSocket,
+        websocket: [
+          connect_info: [:peer_data, :trace_context_headers, :x_headers, :uri, session: [store: :cookie]]
+        ]
+  ```
+
+  With arbitrary keywords:
+
+  ```
+    socket "/socket", AppWeb.UserSocket,
+        websocket: [
+          connect_info: [:uri, custom_value: "abcdef"]
+        ]
+  ```
 
   ## Websocket configuration
 
diff --git a/lib/phoenix/endpoint/supervisor.ex b/lib/phoenix/endpoint/supervisor.ex
index a13fb9479..cd44ffdda 100644
--- a/lib/phoenix/endpoint/supervisor.ex
+++ b/lib/phoenix/endpoint/supervisor.ex
@@ -78,8 +78,9 @@ defmodule Phoenix.Endpoint.Supervisor do
     children =
       config_children(mod, secret_conf, default_conf) ++
         pubsub_children(mod, conf) ++
-        socket_children(mod) ++
+        socket_children(mod, :child_spec) ++
         server_children(mod, conf, server?) ++
+        socket_children(mod, :drainer_spec) ++
         watcher_children(mod, conf, server?)
 
     Supervisor.init(children, strategy: :one_for_one)
@@ -111,10 +112,21 @@ defmodule Phoenix.Endpoint.Supervisor do
     end
   end
 
-  defp socket_children(endpoint) do
-    endpoint.__sockets__
-    |> Enum.uniq_by(&elem(&1, 1))
-    |> Enum.map(fn {_, socket, opts} -> socket.child_spec([endpoint: endpoint] ++ opts) end)
+  defp socket_children(endpoint, fun) do
+    for {_, socket, opts} <- Enum.uniq_by(endpoint.__sockets__, &elem(&1, 1)),
+        spec = apply_or_ignore(socket, fun, [[endpoint: endpoint] ++ opts]),
+        spec != :ignore do
+      spec
+    end
+  end
+
+  defp apply_or_ignore(socket, fun, args) do
+    # If the module is not loaded, we want to invoke and crash
+    if not Code.ensure_loaded?(socket) or function_exported?(socket, fun, length(args)) do
+      apply(socket, fun, args)
+    else
+      :ignore
+    end
   end
 
   defp config_children(mod, conf, default_conf) do
diff --git a/lib/phoenix/socket.ex b/lib/phoenix/socket.ex
index 7de5bfa47..86ba88f7d 100644
--- a/lib/phoenix/socket.ex
+++ b/lib/phoenix/socket.ex
@@ -164,6 +164,19 @@ defmodule Phoenix.Socket do
         {:stop, {:shutdown, :left}, socket}
       end
 
+  A special message delivered to all channels is a Broadcast with
+  event "phx_drain", which is sent when draining the socket during
+  application shutdown. Typically it is handled by sending a drain
+  message to the transport, causing it to shutdown:
+
+      def handle_info(
+            %Broadcast{event: "phx_drain"},
+            %{transport_pid: transport_pid} = socket
+          ) do
+        send(transport_pid, :socket_drain)
+        {:stop, {:shutdown, :draining}, socket}
+      end
+
   We also recommend all channels to monitor the `transport_pid`
   on `init` and exit if the transport exits. We also advise to rewrite
   `:normal` exit reasons (usually due to the socket being closed)
@@ -179,9 +192,7 @@ defmodule Phoenix.Socket do
   a `{:socket_close, pid, reason}` message is sent to the socket before
   shutdown.
 
-  Custom channel implementations cannot be tested with `Phoenix.ChannelTest`
-  and are currently considered experimental. The underlying API may be
-  changed at any moment.
+  Custom channel implementations cannot be tested with `Phoenix.ChannelTest`.
   """
 
   require Logger
@@ -289,6 +300,11 @@ defmodule Phoenix.Socket do
         Phoenix.Socket.__child_spec__(__MODULE__, opts, @phoenix_socket_options)
       end
 
+      @doc false
+      def drainer_spec(opts) do
+        Phoenix.Socket.__drainer_spec__(__MODULE__, opts, @phoenix_socket_options)
+      end
+
       @doc false
       def connect(map), do: Phoenix.Socket.__connect__(__MODULE__, map, @phoenix_socket_options)
 
@@ -424,6 +440,19 @@ defmodule Phoenix.Socket do
     Supervisor.child_spec({Phoenix.Socket.PoolSupervisor, args}, id: handler)
   end
 
+  def __drainer_spec__(handler, opts, socket_options) do
+    endpoint = Keyword.fetch!(opts, :endpoint)
+    opts = Keyword.merge(socket_options, opts)
+
+    if drainer = Keyword.get(opts, :drainer, []) do
+      shutdown = Keyword.get(drainer, :shutdown, 15_000)
+      interval = Keyword.get(drainer, :drain_check_interval, 1_000)
+      {Phoenix.Socket.PoolDrainer, {endpoint, handler, shutdown, interval}}
+    else
+      :ignore
+    end
+  end
+
   def __connect__(user_socket, map, socket_options) do
     %{
       endpoint: endpoint,
@@ -492,12 +521,16 @@ defmodule Phoenix.Socket do
     {:stop, {:shutdown, :disconnected}, state}
   end
 
+  def __info__(:socket_drain, state) do
+    {:stop, {:shutdown, :draining}, state}
+  end
+
   def __info__({:socket_push, opcode, payload}, state) do
     {:push, {opcode, payload}, state}
   end
 
-  def __info__({:socket_close, pid, _reason}, {state, socket}) do
-    socket_close(pid, {state, socket})
+  def __info__({:socket_close, pid, _reason}, state) do
+    socket_close(pid, state)
   end
 
   def __info__(:garbage_collect, state) do
diff --git a/lib/phoenix/socket/pool_supervisor.ex b/lib/phoenix/socket/pool_supervisor.ex
index bf9c1c27d..838c31d71 100644
--- a/lib/phoenix/socket/pool_supervisor.ex
+++ b/lib/phoenix/socket/pool_supervisor.ex
@@ -2,7 +2,6 @@ defmodule Phoenix.Socket.PoolSupervisor do
   @moduledoc false
   use Supervisor
 
-  # TODO: Use PartitionSupervisor once we require Elixir v1.14
   def start_link({endpoint, name, partitions}) do
     Supervisor.start_link(
       __MODULE__,
@@ -11,7 +10,9 @@ defmodule Phoenix.Socket.PoolSupervisor do
     )
   end
 
-  def start_child(endpoint, name, key, spec) do
+  def start_child(socket, key, spec) do
+    %{endpoint: endpoint, handler: name} = socket
+
     case endpoint.config({:socket, name}) do
       ets when not is_nil(ets) ->
         partitions = :ets.lookup_element(ets, :partitions, 2)
@@ -31,7 +32,6 @@ defmodule Phoenix.Socket.PoolSupervisor do
     end
   end
 
-  @doc false
   def start_pooled(ref, i) do
     case DynamicSupervisor.start_link(strategy: :one_for_one) do
       {:ok, pid} ->
@@ -43,8 +43,9 @@ defmodule Phoenix.Socket.PoolSupervisor do
     end
   end
 
-  @doc false
+  @impl true
   def init({endpoint, name, partitions}) do
+    # TODO: Use persisent term on Elixir v1.12+
     ref = :ets.new(name, [:public, read_concurrency: true])
     :ets.insert(ref, {:partitions, partitions})
     Phoenix.Config.permanent(endpoint, {:socket, name}, ref)
@@ -62,3 +63,77 @@ defmodule Phoenix.Socket.PoolSupervisor do
     Supervisor.init(children, strategy: :one_for_one)
   end
 end
+
+defmodule Phoenix.Socket.PoolDrainer do
+  @moduledoc false
+  use GenServer
+  require Logger
+
+  def child_spec({_endpoint, name, shutdown, _interval} = tuple) do
+    # The process should terminate within shutdown but,
+    # in case it doesn't, we will be killed if we exceed
+    # double of that
+    %{
+      id: {:terminator, name},
+      start: {__MODULE__, :start_link, [tuple]},
+      shutdown: shutdown * 2
+    }
+  end
+
+  def start_link(tuple) do
+    GenServer.start_link(__MODULE__, tuple)
+  end
+
+  @impl true
+  def init(tuple) do
+    Process.flag(:trap_exit, true)
+    {:ok, tuple}
+  end
+
+  @impl true
+  def terminate(_reason, {endpoint, name, shutdown, interval}) do
+    ets = endpoint.config({:socket, name})
+    partitions = :ets.lookup_element(ets, :partitions, 2)
+
+    {collection, total} =
+      Enum.map_reduce(0..(partitions - 1), 0, fn index, total ->
+        try do
+          sup = :ets.lookup_element(ets, index, 2)
+          children = DynamicSupervisor.which_children(sup)
+          {Enum.map(children, &elem(&1, 1)), total + length(children)}
+        catch
+          _, _ -> {[], total}
+        end
+      end)
+
+    rounds = div(shutdown, interval)
+    batch = max(ceil(total / rounds), 1)
+
+    if total != 0 do
+      Logger.info("Shutting down #{total} sockets in #{shutdown}ms in #{rounds} rounds")
+    end
+
+    for pids <- collection |> Stream.concat() |> Stream.chunk_every(batch) do
+      {_pid, ref} =
+        spawn_monitor(fn ->
+          refs =
+            for pid <- pids do
+              send(pid, %Phoenix.Socket.Broadcast{event: "phx_draining"})
+              Process.monitor(pid)
+            end
+
+          Enum.each(refs, fn _ ->
+            receive do
+              {:DOWN, _, _, _, _} -> :ok
+            end
+          end)
+        end)
+
+      receive do
+        {:DOWN, ^ref, _, _, _} -> :ok
+      after
+        interval -> Process.demonitor(ref, [:flush])
+      end
+    end
+  end
+end
diff --git a/lib/phoenix/socket/transport.ex b/lib/phoenix/socket/transport.ex
index 24ed391ff..e5a99a7ce 100644
--- a/lib/phoenix/socket/transport.ex
+++ b/lib/phoenix/socket/transport.ex
@@ -21,8 +21,8 @@ defmodule Phoenix.Socket.Transport do
         @behaviour Phoenix.Socket.Transport
 
         def child_spec(opts) do
-          # We won't spawn any process, so let's return a dummy task
-          %{id: __MODULE__, start: {Task, :start_link, [fn -> :ok end]}, restart: :transient}
+          # We won't spawn any process, so let's ignore the child spec
+          :ignore
         end
 
         def connect(state) do
@@ -115,8 +115,22 @@ defmodule Phoenix.Socket.Transport do
       socket "/my_app", MyApp.Socket, shutdown: 5000
 
   means `child_spec([shutdown: 5000])` will be invoked.
+
+  `:ignore` means no child spec is necessary for this socket.
+  """
+  @callback child_spec(keyword) :: :supervisor.child_spec() | :ignore
+
+  @doc """
+  Returns a child specification for terminating the socket.
+
+  This is a process that is started late in the supervision
+  tree with the specific goal of draining connections on
+  application shutdown.
+
+  Similar to `child_spec/1`, it receives the socket options
+  from the endpoint.
   """
-  @callback child_spec(keyword) :: :supervisor.child_spec()
+  @callback drainer_spec(keyword) :: :supervisor.child_spec() | :ignore
 
   @doc """
   Connects to the socket.
@@ -224,7 +238,7 @@ defmodule Phoenix.Socket.Transport do
   """
   @callback terminate(reason :: term, state) :: :ok
 
-  @optional_callbacks handle_control: 2
+  @optional_callbacks handle_control: 2, drainer_spec: 1
 
   require Logger
 
diff --git a/priv/templates/phx.gen.live/core_components.ex b/priv/templates/phx.gen.live/core_components.ex
index fab02c0b4..4823704a5 100644
--- a/priv/templates/phx.gen.live/core_components.ex
+++ b/priv/templates/phx.gen.live/core_components.ex
@@ -254,7 +254,7 @@ defmodule <%= @web_namespace %>.CoreComponents do
   attr :prompt, :string, default: nil, doc: "the prompt for select inputs"
   attr :options, :list, doc: "the options to pass to Phoenix.HTML.Form.options_for_select/2"
   attr :multiple, :boolean, default: false, doc: "the multiple flag for select inputs"
-  attr :rest, :global, include: ~w(autocomplete cols disabled form max maxlength min minlength
+  attr :rest, :global, include: ~w(autocomplete cols disabled form list max maxlength min minlength
                                    pattern placeholder readonly required rows size step)
   slot :inner_block
 


Test patch:
diff --git a/lib/phoenix/test/channel_test.ex b/lib/phoenix/test/channel_test.ex
index 68558af26..da41c2a65 100644
--- a/lib/phoenix/test/channel_test.ex
+++ b/lib/phoenix/test/channel_test.ex
@@ -412,7 +412,7 @@ defmodule Phoenix.ChannelTest do
     %Socket{transport: {__MODULE__, sup}} = socket
 
     starter =
-      fn _, _, _, spec ->
+      fn _, _, spec ->
         Supervisor.start_child(sup, %{spec | id: make_ref()})
       end
 
diff --git a/test/phoenix/integration/long_poll_socket_test.exs b/test/phoenix/integration/long_poll_socket_test.exs
index 34e236aa0..47357f66c 100644
--- a/test/phoenix/integration/long_poll_socket_test.exs
+++ b/test/phoenix/integration/long_poll_socket_test.exs
@@ -27,7 +27,7 @@ defmodule Phoenix.Integration.LongPollSocketTest do
 
     def child_spec(opts) do
       :value = Keyword.fetch!(opts, :custom)
-      Supervisor.child_spec({Task, fn -> :ok end}, [])
+      :ignore
     end
 
     def connect(map) do
diff --git a/test/phoenix/integration/websocket_socket_test.exs b/test/phoenix/integration/websocket_socket_test.exs
index ad99c9480..5bc4c5599 100644
--- a/test/phoenix/integration/websocket_socket_test.exs
+++ b/test/phoenix/integration/websocket_socket_test.exs
@@ -25,7 +25,7 @@ defmodule Phoenix.Integration.WebSocketTest do
 
     def child_spec(opts) do
       :value = Keyword.fetch!(opts, :custom)
-      Supervisor.child_spec({Task, fn -> :ok end}, [])
+      :ignore
     end
 
     def connect(map) do
@@ -60,7 +60,7 @@ defmodule Phoenix.Integration.WebSocketTest do
   defmodule PingSocket do
     @behaviour Phoenix.Socket.Transport
 
-    def child_spec(_opts), do: Supervisor.child_spec({Task, fn -> :ok end}, [id: :ping])
+    def child_spec(_opts), do: :ignore
     def connect(_), do: {:ok, %{}}
     def init(state), do: {:ok, state}