Allow a leaving channel to shut down gracefully

lib/phoenix/socket.ex

@@ -769,19 +769,32 @@ defmodule Phoenix.Socket do
   end
 
   defp handle_in({pid, _ref, status}, %{event: "phx_join", topic: topic} = message, state, socket) do
-    receive do
-      {:socket_close, ^pid, _reason} -> :ok
-    after
-      0 ->
-        if status != :leaving do
-          Logger.debug(fn ->
-            "Duplicate channel join for topic \"#{topic}\" in #{inspect(socket.handler)}. " <>
-              "Closing existing channel for new join."
-          end)
-        end
-    end
+    closed? =
+      receive do
+        {:socket_close, ^pid, _reason} -> true
+      after
+        0 ->
+          if status != :leaving do
+            Logger.debug(fn ->
+              "Duplicate channel join for topic \"#{topic}\" in #{inspect(socket.handler)}. " <>
+                "Closing existing channel for new join."
+            end)
+          end
+
+          false
+      end
+
+    :ok =
+      if status == :leaving and not closed? do
+        # the channel is already leaving on the client's request, so it stops on
+        # its own once it handled the messages the client sent before the leave.
+        # We wait for it instead of shutting it down, as shutting it down would
+        # discard those messages
+        await_channel_shutdown(pid)
+      else
+        shutdown_duplicate_channel(pid)
+      end
 
-    :ok = shutdown_duplicate_channel(pid)
     {:push, {opcode, payload}, {new_state, new_socket}} = socket_close(pid, {state, socket})
     send(self(), {:socket_push, opcode, payload})
     handle_in(nil, message, new_state, new_socket)
@@ -912,6 +925,18 @@ defmodule Phoenix.Socket do
     end
   end
 
+  defp await_channel_shutdown(pid) do
+    ref = Process.monitor(pid)
+
+    receive do
+      {:DOWN, ^ref, _, _, _} -> :ok
+    after
+      5_000 ->
+        Process.exit(pid, :kill)
+        receive do: ({:DOWN, ^ref, _, _, _} -> :ok)
+    end
+  end
+
   defp socket_close(pid, {state, socket}) do
     case state.channels_inverse do
       %{^pid => {topic, join_ref}} ->

test/phoenix/integration/websocket_channels_test.exs

@@ -49,6 +49,11 @@ defmodule Phoenix.Integration.WebSocketChannelsTest do
       {:reply, :ok, socket}
     end
 
+    def handle_in("slow_msg", _message, socket) do
+      Process.sleep(100)
+      {:reply, :ok, socket}
+    end
+
     def handle_in("boom", _message, _socket) do
       raise "boom"
     end
@@ -820,6 +825,45 @@ defmodule Phoenix.Integration.WebSocketChannelsTest do
           }
         end
 
+        test "rejoining a leaving topic waits for the channel to handle pending messages" do
+          {:ok, sock} = WebsocketClient.connect(self(), "#{@vsn_path}&user_id=1001", @serializer)
+          WebsocketClient.join(sock, "room:rejoin", %{})
+
+          assert_receive %Message{
+            topic: "room:rejoin",
+            event: "phx_reply",
+            ref: "1",
+            payload: %{"response" => %{}, "status" => "ok"}
+          }
+
+          chan = Process.whereis(:"room:rejoin")
+          Process.monitor(chan)
+
+          # the client pushes an event and then immediately leaves and rejoins
+          # the same topic, as the LiveView client does on live navigation;
+          # the pushed event must still be handled
+          WebsocketClient.send_event(sock, "room:rejoin", "slow_msg", %{})
+          WebsocketClient.leave(sock, "room:rejoin", %{})
+          WebsocketClient.join(sock, "room:rejoin", %{})
+
+          assert_receive %Message{
+                           topic: "room:rejoin",
+                           event: "phx_reply",
+                           ref: "2",
+                           payload: %{"response" => %{}, "status" => "ok"}
+                         },
+                         1000
+
+          assert_receive {:DOWN, _, :process, ^chan, {:shutdown, :left}}, 1000
+
+          assert_receive %Message{
+            topic: "room:rejoin",
+            event: "phx_reply",
+            ref: "4",
+            payload: %{"response" => %{}, "status" => "ok"}
+          }
+        end
+
         test "returns 403 when versions to not match" do
           assert capture_log(fn ->
                    url = "ws://127.0.0.1:#{@port}/ws/websocket?vsn=123.1.1"