From ab39c680e2485b3948461fa972e36f849feb8f47 Mon Sep 17 00:00:00 2001
From: Chris McCord <chris@chrismccord.com>
Date: Thu, 10 Aug 2023 23:53:04 -0400
Subject: [PATCH] AsyncResult WIP

---
 lib/phoenix_live_view.ex              |  35 ++
 lib/phoenix_live_view/async_assign.ex | 314 ------------------
 lib/phoenix_live_view/async_result.ex | 449 ++++++++++++++++++++++++++
 lib/phoenix_live_view/channel.ex      |  28 +-
 test/phoenix_live_view_test.exs       |  13 +
 test/support/live_views/general.ex    |  37 ++-
 6 files changed, 532 insertions(+), 344 deletions(-)
 delete mode 100644 lib/phoenix_live_view/async_assign.ex
 create mode 100644 lib/phoenix_live_view/async_result.ex

diff --git a/lib/phoenix_live_view.ex b/lib/phoenix_live_view.ex
index 4ef2fe5223..c96ef37b2f 100644
--- a/lib/phoenix_live_view.ex
+++ b/lib/phoenix_live_view.ex
@@ -666,6 +666,40 @@ defmodule Phoenix.LiveView do
   """
   def connected?(%Socket{transport_pid: transport_pid}), do: transport_pid != nil
 
+  @doc """
+  Puts a new private key and value in the socket.
+
+  Privates are *not change tracked*. This storage is meant to be used by
+  users and libraries to hold state that doesn't require
+  change tracking. The keys should be prefixed with the app/library name.
+
+  ## Examples
+
+  Key values can be placed in private:
+
+      put_private(socket, :myapp_meta, %{foo: "bar"})
+
+  And then retreived:
+
+      socket.private[:myapp_meta]
+  """
+  @reserved_privates ~w(
+    connect_params
+    connect_info
+    assign_new
+    live_layout
+    lifecycle
+    root_view
+    __temp__
+  )a
+  def put_private(%Socket{} = socket, key, value) when key not in @reserved_privates do
+    %Socket{socket | private: Map.put(socket.private, key, value)}
+  end
+
+  def put_private(%Socket{}, bad_key, _value) do
+    raise ArgumentError, "cannot set reserved private key #{inspect(bad_key)}"
+  end
+
   @doc """
   Adds a flash message to the socket to be displayed.
 
@@ -690,6 +724,7 @@ defmodule Phoenix.LiveView do
       iex> put_flash(socket, :info, "It worked!")
       iex> put_flash(socket, :error, "You can't access that page")
   """
+
   defdelegate put_flash(socket, kind, msg), to: Phoenix.LiveView.Utils
 
   @doc """
diff --git a/lib/phoenix_live_view/async_assign.ex b/lib/phoenix_live_view/async_assign.ex
deleted file mode 100644
index c5d67c2987..0000000000
--- a/lib/phoenix_live_view/async_assign.ex
+++ /dev/null
@@ -1,314 +0,0 @@
-defmodule Phoenix.LiveView.AsyncAssign do
-  @moduledoc ~S'''
-  Adds async_assign functionality to LiveViews and LiveComponents.
-
-  Performing asynchronous work is common in LiveViews and LiveComponents.
-  It allows the user to get a working UI quicky while the system fetches some
-  data in the background or talks to an external service. For async work,
-  you also typically need to handle the different states of the async operation,
-  such as loading, error, and the successful result. You also want to catch any
-  error and translate it to a meaningful update in the UI rather than crashing
-  the user experience.
-
-  ## Examples
-
-  The `assign_async/3` function takes a name, a list of keys which will be assigned
-  asynchronously, and a function that returns the result of the async operation.
-  For example, let's say we want to async fetch a user's organization from the database,
-  as well as their profile and rank:
-
-      def mount(%{"slug" => slug}, _, socket) do
-        {:ok,
-         socket
-         |> assign(:foo, "bar")
-         |> assign_async(:org, fn -> {:ok, %{org: fetch_org!(slug)} end)
-         |> assign_async(:profile, [:profile, :rank], fn -> {:ok, %{profile: ..., rank: ...}} end)}
-      end
-
-  Here we are assigning `:org` and `[:profile, :rank]` asynchronously. If no keys are
-  given (as in the case of `:org`), the keys will default to `[:org]`. The async function
-  must return a `{:ok, assigns}` or `{:error, reason}` tuple where `assigns` is a map of
-  the keys passed to `assign_async`. If the function returns other keys or a different
-  set of keys, an error is raised.
-
-  The state of the async operation is stored in the socket assigns under the `@async` assign
-  on the socket with the name given to `assign_async/3`. It carries the `:loading?`,
-  `:error`, and `:result` keys. For example, if we wanted to show the loading states
-  in the UI for the `:org`, our template could conditionally render the states:
-
-  ```heex
-  <div :if={@async.org.loading?}>Loading organization...</div>
-  <div :if={@async.org.error}>there was an error loading the organization</div>
-  <div :if={@async.org.result == nil}}>You don't have an org yet</div>
-  <div :if={org = @async.org.result}}><%= org.name%> loaded!</div>
-  ```
-
-  The `async_result` function component can also be used to declaratively
-  render the different states using slots:
-
-  ```heex
-  <.async_result :let={org} assign={@async.org}>
-    <:loading>Loading organization...</:loading>
-    <:empty>You don't have an organization yet</:error>
-    <:error>there was an error loading the organization</:error>
-    <%= org.name %>
-  <.async_result>
-  ```
-
-  Additionally, for async assigns which result in a list of items, you
-  can consume the `@async.<name>` directly, and it will only enumerate
-  the results once the results are loaded. For example:
-
-  ```heex
-  <div :for={orgs <- @async.orgs}><%= org.name %></div>
-  ```
-  '''
-  use Phoenix.Component
-
-  defstruct name: nil,
-            ref: nil,
-            pid: nil,
-            keys: [],
-            loading?: false,
-            error: nil,
-            result: nil,
-            canceled?: false
-
-  alias Phoenix.LiveView.AsyncAssign
-
-  @doc """
-  Renders an async assign with slots for the different loading states.
-
-  ## Examples
-
-  ```heex
-  <.async_result :let={org} assign={@async.org}>
-    <:loading>Loading organization...</:loading>
-    <:empty>You don't have an organization yet</:error>
-    <:error :let={_reason}>there was an error loading the organization</:error>
-    <:canceled :let={_reason}>loading cancled</:canceled>
-    <%= org.name %>
-  <.async_result>
-  ```
-  """
-  attr :assign, :any, required: true
-  slot :loading
-  slot :canceled
-
-  slot :empty,
-    doc:
-      "rendered when the result is loaded and is either nil or an empty list. Receives the result as a :let."
-
-  slot :error,
-    doc:
-      "rendered when an error is caught or the function return `{:error, reason}`. Receives the error as a :let."
-
-  def async_result(assigns) do
-    case assigns.assign do
-      %AsyncAssign{result: result, loading?: false, error: nil, canceled?: false} ->
-        if assigns.empty != [] && result in [nil, []] do
-          ~H|<%= render_slot(@empty, @assign.result) %>|
-        else
-          ~H|<%= render_slot(@inner_block, @assign.result) %>|
-        end
-
-      %AsyncAssign{loading?: true} ->
-        ~H|<%= render_slot(@loading) %>|
-
-      %AsyncAssign{loading?: false, error: error} when not is_nil(error) ->
-        ~H|<%= render_slot(@error, @assign.error) %>|
-
-      %AsyncAssign{loading?: false, canceled?: true} ->
-        ~H|<%= render_slot(@canceled) %>|
-    end
-  end
-
-  @doc """
-  Assigns keys ansynchronously.
-
-  See the module docs for more and exmaple usage.
-  """
-  def assign_async(%Phoenix.LiveView.Socket{} = socket, name, func) do
-    assign_async(socket, name, [name], func)
-  end
-
-  def assign_async(%Phoenix.LiveView.Socket{} = socket, name, keys, func)
-      when is_atom(name) and is_list(keys) and is_function(func, 0) do
-    socket = cancel_existing(socket, name)
-    base_async = %AsyncAssign{name: name, keys: keys, loading?: true}
-
-    async_assign =
-      if Phoenix.LiveView.connected?(socket) do
-        lv_pid = self()
-        ref = make_ref()
-        cid = if myself = socket.assigns[:myself], do: myself.cid
-
-        {:ok, pid} =
-          Task.start_link(fn ->
-            do_async(lv_pid, cid, %AsyncAssign{base_async | pid: self(), ref: ref}, func)
-          end)
-
-        %AsyncAssign{base_async | pid: pid, ref: ref}
-      else
-        base_async
-      end
-
-    Enum.reduce(keys, socket, fn key, acc -> update_async(acc, key, async_assign) end)
-  end
-
-  @doc """
-  Cancels an async assign.
-
-  ## Examples
-
-      def handle_event("cancel_preview", _, socket) do
-        {:noreply, cancel_async(socket, :preview)}
-      end
-  """
-  def cancel_async(%Phoenix.LiveView.Socket{} = socket, name) do
-    case get(socket, name) do
-      %AsyncAssign{loading?: false} ->
-        socket
-
-      %AsyncAssign{canceled?: false, pid: pid} = existing ->
-        Process.unlink(pid)
-        Process.exit(pid, :kill)
-        async = %AsyncAssign{existing | loading?: false, error: nil, canceled?: true}
-        update_async(socket, name, async)
-
-      nil ->
-        raise ArgumentError,
-              "no async assign #{inspect(name)} previously assigned with assign_async"
-    end
-  end
-
-  defp do_async(lv_pid, cid, %AsyncAssign{} = async_assign, func) do
-    %AsyncAssign{keys: known_keys, name: name, ref: ref} = async_assign
-
-    try do
-      case func.() do
-        {:ok, %{} = results} ->
-          if Map.keys(results) -- known_keys == [] do
-            Phoenix.LiveView.Channel.write_socket(lv_pid, cid, fn socket ->
-              write_current_async(socket, name, ref, fn %AsyncAssign{} = current ->
-                Enum.reduce(results, socket, fn {key, result}, acc ->
-                  async = %AsyncAssign{current | result: result, loading?: false}
-                  update_async(acc, key, async)
-                end)
-              end)
-            end)
-          else
-            raise ArgumentError, """
-            expected assign_async to return map of
-            assigns for all keys in #{inspect(known_keys)}, but got: #{inspect(results)}
-            """
-          end
-
-        {:error, reason} ->
-          Phoenix.LiveView.Channel.write_socket(lv_pid, cid, fn socket ->
-            write_current_async(socket, name, ref, fn %AsyncAssign{} = current ->
-              Enum.reduce(known_keys, socket, fn key, acc ->
-                async = %AsyncAssign{current | result: nil, loading?: false, error: reason}
-                update_async(acc, key, async)
-              end)
-            end)
-          end)
-
-        other ->
-          raise ArgumentError, """
-          expected assign_async to return {:ok, map} of
-          assigns for #{inspect(known_keys)} or {:error, reason}, got: #{inspect(other)}
-          """
-      end
-    catch
-      kind, reason ->
-        Process.unlink(lv_pid)
-
-        Phoenix.LiveView.Channel.write_socket(lv_pid, cid, fn socket ->
-          write_current_async(socket, name, ref, fn %AsyncAssign{} = current ->
-            Enum.reduce(known_keys, socket, fn key, acc ->
-              async = %AsyncAssign{current | loading?: false, error: {kind, reason}}
-              update_async(acc, key, async)
-            end)
-          end)
-        end)
-
-        :erlang.raise(kind, reason, __STACKTRACE__)
-    end
-  end
-
-  defp update_async(socket, key, %AsyncAssign{} = new_async) do
-    socket
-    |> ensure_async()
-    |> Phoenix.Component.update(:async, fn async_map ->
-      Map.put(async_map, key, new_async)
-    end)
-  end
-
-  defp ensure_async(socket) do
-    Phoenix.Component.assign_new(socket, :async, fn -> %{} end)
-  end
-
-  defp get(%Phoenix.LiveView.Socket{} = socket, name) do
-    socket.assigns[:async] && Map.get(socket.assigns.async, name)
-  end
-
-  # handle race of async being canceled and then reassigned
-  defp write_current_async(socket, name, ref, func) do
-    case get(socket, name) do
-      %AsyncAssign{ref: ^ref} = async_assign -> func.(async_assign)
-      %AsyncAssign{ref: _ref} -> socket
-    end
-  end
-
-  defp cancel_existing(socket, name) do
-    if get(socket, name) do
-      cancel_async(socket, name)
-    else
-      socket
-    end
-  end
-
-  defimpl Enumerable, for: Phoenix.LiveView.AsyncAssign do
-    alias Phoenix.LiveView.AsyncAssign
-
-    def count(%AsyncAssign{result: result, loading?: false, error: nil, canceled?: false}),
-      do: Enum.count(result)
-
-    def count(%AsyncAssign{}), do: 0
-
-    def member?(%AsyncAssign{result: result, loading?: false, error: nil, canceled?: false}, item) do
-      Enum.member?(result, item)
-    end
-
-    def member?(%AsyncAssign{}, _item) do
-      raise RuntimeError, "cannot lookup member? while loading"
-    end
-
-    def reduce(
-          %AsyncAssign{result: result, loading?: false, error: nil, canceled?: false},
-          acc,
-          fun
-        ) do
-      do_reduce(result, acc, fun)
-    end
-
-    def reduce(%AsyncAssign{}, acc, _fun), do: acc
-
-    defp do_reduce(_list, {:halt, acc}, _fun), do: {:halted, acc}
-    defp do_reduce(list, {:suspend, acc}, fun), do: {:suspended, acc, &do_reduce(list, &1, fun)}
-    defp do_reduce([], {:cont, acc}, _fun), do: {:done, acc}
-
-    defp do_reduce([item | tail], {:cont, acc}, fun) do
-      do_reduce(tail, fun.(item, acc), fun)
-    end
-
-    def slice(%AsyncAssign{result: result, loading?: false, error: nil, canceled?: false}) do
-      fn start, length, step -> Enum.slice(result, start..(start + length - 1)//step) end
-    end
-
-    def slice(%AsyncAssign{}) do
-      fn _start, _length, _step -> [] end
-    end
-  end
-end
diff --git a/lib/phoenix_live_view/async_result.ex b/lib/phoenix_live_view/async_result.ex
new file mode 100644
index 0000000000..bed8c27cf2
--- /dev/null
+++ b/lib/phoenix_live_view/async_result.ex
@@ -0,0 +1,449 @@
+defmodule Phoenix.LiveView.AsyncResult do
+  @moduledoc ~S'''
+  Adds async_assign functionality to LiveViews and LiveComponents.
+
+  Performing asynchronous work is common in LiveViews and LiveComponents.
+  It allows the user to get a working UI quicky while the system fetches some
+  data in the background or talks to an external service. For async work,
+  you also typically need to handle the different states of the async operation,
+  such as loading, error, and the successful result. You also want to catch any
+  error and translate it to a meaningful update in the UI rather than crashing
+  the user experience.
+
+  ## Examples
+
+  The `assign_async/3` function takes a name, a list of keys which will be assigned
+  asynchronously, and a function that returns the result of the async operation.
+  For example, let's say we want to async fetch a user's organization from the database,
+  as well as their profile and rank:
+
+      def mount(%{"slug" => slug}, _, socket) do
+        {:ok,
+         socket
+         |> assign(:foo, "bar")
+         |> assign_async(:org, fn -> {:ok, %{org: fetch_org!(slug)} end)
+         |> assign_async(:profile, [:profile, :rank], fn -> {:ok, %{profile: ..., rank: ...}} end)}
+      end
+
+  Here we are assigning `:org` and `[:profile, :rank]` asynchronously. If no keys are
+  given (as in the case of `:org`), the keys will default to `[:org]`. The async function
+  must return a `{:ok, assigns}` or `{:error, reason}` tuple where `assigns` is a map of
+  the keys passed to `assign_async`. If the function returns other keys or a different
+  set of keys, an error is raised.
+
+  The state of the async operation is stored in the socket assigns under the `@async` assign
+  on the socket with the name given to `assign_async/3`. It carries the `:loading?`,
+  `:error`, and `:result` keys. For example, if we wanted to show the loading states
+  in the UI for the `:org`, our template could conditionally render the states:
+
+  ```heex
+  <div :if={@async.org.loading?}>Loading organization...</div>
+  <div :if={@async.org.error}>there was an error loading the organization</div>
+  <div :if={@async.org.result == nil}}>You don't have an org yet</div>
+  <div :if={org = @async.org.result}}><%= org.name%> loaded!</div>
+  ```
+
+  The `async_result` function component can also be used to declaratively
+  render the different states using slots:
+
+  ```heex
+  <.async_result :let={org} assign={@async.org}>
+    <:loading>Loading organization...</:loading>
+    <:empty>You don't have an organization yet</:error>
+    <:error>there was an error loading the organization</:error>
+    <%= org.name %>
+  <.async_result>
+  ```
+
+  Additionally, for async assigns which result in a list of items, you
+  can consume the `@async.<name>` directly, and it will only enumerate
+  the results once the results are loaded. For example:
+
+  ```heex
+  <div :for={orgs <- @async.orgs}><%= org.name %></div>
+  ```
+  '''
+  use Phoenix.Component
+
+  defstruct name: nil,
+            keys: [],
+            ok?: false,
+            state: :loading,
+            result: nil
+
+  alias Phoenix.LiveView.{AsyncResult, Socket}
+
+  @doc """
+  TODO
+  """
+  def new(name, keys) do
+    loading(%AsyncResult{name: name, keys: keys, result: nil, ok?: false})
+  end
+
+  @doc """
+  TODO
+  """
+  def loading(%AsyncResult{} = result) do
+    %AsyncResult{result | state: :loading}
+  end
+
+  @doc """
+  TODO
+  """
+  def canceled(%AsyncResult{} = result) do
+    %AsyncResult{result | state: :canceled}
+  end
+
+  @doc """
+  TODO
+  """
+  def error(%AsyncResult{} = result, reason) do
+    %AsyncResult{result | state: {:error, reason}}
+  end
+
+  @doc """
+  TODO
+  """
+  def exit(%AsyncResult{} = result, reason) do
+    %AsyncResult{result | state: {:exit, reason}}
+  end
+
+  @doc """
+  TODO
+  """
+  def throw(%AsyncResult{} = result, value) do
+    %AsyncResult{result | state: {:throw, value}}
+  end
+
+  @doc """
+  TODO
+  """
+  def ok(%AsyncResult{} = result, value) do
+    %AsyncResult{result | state: :ok, ok?: true, result: value}
+  end
+
+  @doc """
+  Renders an async assign with slots for the different loading states.
+
+  ## Examples
+
+  ```heex
+  <AsyncResult.with_state :let={org} assign={@org}>
+    <:loading>Loading organization...</:loading>
+    <:empty>You don't have an organization yet</:error>
+    <:error :let={_reason}>there was an error loading the organization</:error>
+    <:canceled :let={_reason}>loading cancled</:canceled>
+    <%= org.name %>
+  <AsyncResult.with_state>
+  ```
+  """
+  attr :assign, :any, required: true
+  slot :loading
+  slot :canceled
+
+  slot :empty,
+    doc:
+      "rendered when the result is loaded and is either nil or an empty list. Receives the result as a :let."
+
+  slot :error,
+    doc:
+      "rendered when an error is caught or the function return `{:error, reason}`. Receives the error as a :let."
+
+  def with_state(assigns) do
+    case assigns.assign do
+      %AsyncResult{state: :ok, result: result} ->
+        if assigns.empty != [] && result in [nil, []] do
+          ~H|<%= render_slot(@empty, @assign.result) %>|
+        else
+          ~H|<%= render_slot(@inner_block, @assign.result) %>|
+        end
+
+      %AsyncResult{state: :loading} ->
+        ~H|<%= render_slot(@loading) %>|
+
+      %AsyncResult{state: :canceled} ->
+        ~H|<%= render_slot(@canceled) %>|
+
+      %AsyncResult{state: {kind, _value}} when kind in [:error, :exit, :throw] ->
+        ~H|<%= render_slot(@error, @assign.state) %>|
+    end
+  end
+
+  @doc """
+  Assigns keys asynchronously.
+
+  The task is linked to the caller and errors are wrapped.
+  Each key passed to `assign_async/3` will be assigned to
+  an `%AsyncResult{}` struct holding the status of the operation
+  and the result when completed.
+  """
+  def assign_async(%Socket{} = socket, key_or_keys, func)
+      when (is_atom(key_or_keys) or is_list(key_or_keys)) and
+             is_function(func, 0) do
+    keys = List.wrap(key_or_keys)
+
+    keys
+    |> Enum.reduce(socket, fn key, acc ->
+      Phoenix.Component.assign(acc, key, AsyncResult.new(key, keys))
+    end)
+    |> run_async_task(keys, func, fn new_socket, _component_mod, result ->
+      assign_result(new_socket, keys, result)
+    end)
+  end
+
+  defp assign_result(socket, keys, result) do
+    case result do
+      {:ok, %{} = values} ->
+        if Map.keys(values) -- keys == [] do
+          Enum.reduce(values, socket, fn {key, val}, acc ->
+            current_async = get_current_async!(acc, key)
+            Phoenix.Component.assign(acc, key, AsyncResult.ok(current_async, val))
+          end)
+        else
+          raise ArgumentError, """
+          expected assign_async to return map of assigns for all keys
+          in #{inspect(keys)}, but got: #{inspect(values)}
+          """
+        end
+
+      {:error, reason} ->
+        Enum.reduce(keys, socket, fn key, acc ->
+          current_async = get_current_async!(acc, key)
+          Phoenix.Component.assign(acc, key, AsyncResult.error(current_async, reason))
+        end)
+
+      {:exit, reason} ->
+        Enum.reduce(keys, socket, fn key, acc ->
+          current_async = get_current_async!(acc, key)
+          Phoenix.Component.assign(acc, key, AsyncResult.exit(current_async, reason))
+        end)
+
+      {:throw, value} ->
+        Enum.reduce(keys, socket, fn key, acc ->
+          current_async = get_current_async!(acc, key)
+          Phoenix.Component.assign(acc, key, AsyncResult.throw(current_async, value))
+        end)
+
+      other ->
+        raise ArgumentError, """
+        expected assign_async to return {:ok, map} of
+        assigns for #{inspect(keys)} or {:error, reason}, got: #{inspect(other)}
+        """
+    end
+  end
+
+  defp get_current_async!(socket, key) do
+    # handle case where assign is temporary and needs to be rebuilt
+    case socket.assigns do
+      %{^key => %AsynResult{} = current_async} -> current_async
+      %{^key => _other} -> AsyncResult.new(key, keys)
+      %{} -> raise ArgumentError, "missing async assign #{inspect(key)}"
+    end
+  end
+
+  @doc """
+  Starts an ansynchronous task.
+
+  The task is linked to the caller and errors are wrapped.
+  The result of the task is sent to the `handle_async/3` callback
+  of the caller LiveView or LiveComponent.
+
+  ## Examples
+
+
+  """
+  def start_async(%Socket{} = socket, name, func)
+      when is_atom(name) and is_function(func, 0) do
+    run_async_task(socket, [name], func, fn new_socket, component_mod, result ->
+      callback_mod = component_mod || new_socket.view
+
+      case result do
+        {tag, value} when tag in [:ok, :error, :exit, :throw] ->
+          :ok
+
+        other ->
+          raise ArgumentError, """
+          expected start_async for #{inspect(name)} in #{inspect(callback_mod)}
+          to return {:ok, result} | {:error, reason}, got:
+
+              #{inspect(other)}
+          """
+      end
+
+      case callback_mod.handle_async(name, result, new_socket) do
+        {:noreply, %Socket{} = new_socket} ->
+          new_socket
+
+        other ->
+          raise ArgumentError, """
+          expected #{inspect(callback_mod)}.handle_async/3 to return {:noreply, socket}, got:
+
+              #{inspect(other)}
+          """
+      end
+    end)
+  end
+
+  defp run_async_task(%Socket{} = socket, keys, func, result_func)
+       when is_list(keys) and is_function(result_func, 3) do
+    if Phoenix.LiveView.connected?(socket) do
+      socket = cancel_existing(socket, keys)
+      lv_pid = self()
+      cid = cid(socket)
+      ref = make_ref()
+      {:ok, pid} = Task.start_link(fn -> do_async(lv_pid, cid, ref, keys, func, result_func) end)
+      update_private_async(socket, keys, {ref, pid})
+    else
+      socket
+    end
+  end
+
+  defp do_async(lv_pid, cid, ref, keys, func, result_func) do
+    try do
+      result = func.()
+
+      Phoenix.LiveView.Channel.write_socket(lv_pid, cid, fn socket, component_mod ->
+        handle_current_async(socket, keys, ref, component_mod, result, result_func)
+      end)
+    catch
+      kind, reason ->
+        Process.unlink(lv_pid)
+
+        Phoenix.LiveView.Channel.write_socket(lv_pid, cid, fn socket, component_mod ->
+          handle_current_async(socket, keys, ref, component_mod, {kind, reason}, result_func)
+        end)
+
+        :erlang.raise(kind, reason, __STACKTRACE__)
+    end
+  end
+
+  # handle race of async being canceled and then reassigned
+  defp handle_current_async(socket, keys, ref, component_mod, result, result_func)
+       when is_function(result_func, 3) do
+    get_private_async(socket, keys) do
+      {^ref, pid} ->
+        new_socket = delete_private_async(socket, keys)
+        result_func.(new_socket, component_mod, result)
+
+      {_ref, _pid} ->
+        socket
+
+      nil ->
+        socket
+    end
+  end
+
+  @doc """
+  Cancels an async assign.
+
+  ## Examples
+
+      TODO fix docs
+      def handle_event("cancel_preview", _, socket) do
+        {:noreply, cancel_async(socket, :preview)}
+      end
+  """
+  def cancel_async(%Socket{} = socket, %AsyncResult{} = result) do
+    result.keys
+    |> Enum.reduce(socket, fn key, acc ->
+      Phoenix.Component.assign(acc, key, AsyncResult.canceled(result))
+    end)
+    |> cancel_async(result.keys)
+  end
+
+  def cancel_async(%Socket{} = socket, key) when is_atom(key) do
+    cancel_async(socket, [key])
+  end
+
+  def cancel_async(%Socket{} = socket, keys) when is_list(keys) do
+    case get_private_async(socket, keys) do
+      {ref, pid} when is_pid(pid) ->
+        Process.unlink(pid)
+        Process.exit(pid, :kill)
+        delete_private_async(socket, keys)
+
+      nil ->
+        raise ArgumentError, "uknown async assign #{inspect(keys)}"
+    end
+  end
+
+  defp update_private_async(socket, keys, {ref, pid}) do
+    socket
+    |> ensure_private_async()
+    |> Phoenix.Component.update(:async, fn async_map ->
+      Map.put(async_map, keys, {ref, pid})
+    end)
+  end
+
+  defp delete_private_async(socket, keys) do
+    socket
+    |> ensure_private_async()
+    |> Phoenix.Component.update(:async, fn async_map -> Map.delete(async_map, keys) end)
+  end
+
+  defp ensure_private_async(socket) do
+    case socket.private do
+      %{phoenix_async: _} -> socket
+      %{} -> Phoenix.LiveView.put_private(socket, :phoenix_async, %{})
+    end
+  end
+
+  defp get_private_async(%Socket{} = socket, keys) do
+    socket.private[:phoenix_async][keys]
+  end
+
+  defp cancel_existing(socket, keys) when is_list(keys) do
+    if get_private_async(acc, keys) do
+      cancel_async(acc, keys)
+    else
+      acc
+    end
+  end
+
+  defp cid(%Socket{} = socket) do
+    if myself = socket.assigns[:myself], do: myself.cid
+  end
+
+  defimpl Enumerable, for: Phoenix.LiveView.AsyncResult do
+    alias Phoenix.LiveView.AsyncResult
+
+    def count(%AsyncResult{result: result, state: :ok}),
+      do: Enum.count(result)
+
+    def count(%AsyncResult{}), do: 0
+
+    def member?(%AsyncResult{result: result, state: :ok}, item) do
+      Enum.member?(result, item)
+    end
+
+    def member?(%AsyncResult{}, _item) do
+      raise RuntimeError, "cannot lookup member? without an ok result"
+    end
+
+    def reduce(
+          %AsyncResult{result: result, state: :ok},
+          acc,
+          fun
+        ) do
+      do_reduce(result, acc, fun)
+    end
+
+    def reduce(%AsyncResult{}, acc, _fun), do: acc
+
+    defp do_reduce(_list, {:halt, acc}, _fun), do: {:halted, acc}
+    defp do_reduce(list, {:suspend, acc}, fun), do: {:suspended, acc, &do_reduce(list, &1, fun)}
+    defp do_reduce([], {:cont, acc}, _fun), do: {:done, acc}
+
+    defp do_reduce([item | tail], {:cont, acc}, fun) do
+      do_reduce(tail, fun.(item, acc), fun)
+    end
+
+    def slice(%AsyncResult{result: result, state: :ok}) do
+      fn start, length, step -> Enum.slice(result, start..(start + length - 1)//step) end
+    end
+
+    def slice(%AsyncResult{}) do
+      fn _start, _length, _step -> [] end
+    end
+  end
+end
diff --git a/lib/phoenix_live_view/channel.ex b/lib/phoenix_live_view/channel.ex
index 5728659a44..ad55707c02 100644
--- a/lib/phoenix_live_view/channel.ex
+++ b/lib/phoenix_live_view/channel.ex
@@ -30,7 +30,7 @@ defmodule Phoenix.LiveView.Channel do
     )
   end
 
-  def write_socket(lv_pid, cid, func) when is_function(func, 1) do
+  def write_socket(lv_pid, cid, func) when is_function(func, 2) do
     GenServer.call(lv_pid, {@prefix, :write_socket, cid, func})
   end
 
@@ -292,23 +292,20 @@ defmodule Phoenix.LiveView.Channel do
 
   def handle_call({@prefix, :async_pids}, _from, state) do
     %{socket: socket} = state
-    lv_pids = get_async_pids(socket.assigns)
+    lv_pids = get_async_pids(socket.private)
 
     component_pids =
       state
-      |> component_assigns()
-      |> Enum.flat_map(fn
-        {_cid, %{async: %{}} = assigns} -> get_async_pids(assigns)
-        {_cid, %{}} -> []
-      end)
+      |> component_privates()
+      |> get_async_pids()
 
     {:reply, {:ok, lv_pids ++ component_pids}, state}
   end
 
   def handle_call({@prefix, :write_socket, cid, func}, _from, state) do
     new_state =
-      write_socket(state, cid, nil, fn socket, _ ->
-        %Phoenix.LiveView.Socket{} = new_socket = func.(socket)
+      write_socket(state, cid, nil, fn socket, maybe_component ->
+        %Phoenix.LiveView.Socket{} = new_socket = func.(socket, maybe_component)
         {new_socket, {:ok, nil, state}}
       end)
 
@@ -1425,15 +1422,18 @@ defmodule Phoenix.LiveView.Channel do
 
   defp maybe_subscribe_to_live_reload(response), do: response
 
-  defp component_assigns(state) do
+  defp component_privates(state) do
     %{components: {components, _ids, _}} = state
 
-    Enum.into(components, %{}, fn {cid, {_mod, _id, assigns, _private, _prints}} ->
-      {cid, assigns}
+    Enum.into(components, %{}, fn {cid, {_mod, _id, _assigns, private, _prints}} ->
+      {cid, private}
     end)
   end
 
-  defp get_async_pids(assigns) do
-    Enum.map(assigns[:async] || %{}, fn {_, %Phoenix.LiveView.AsyncAssign{pid: pid}} -> pid end)
+  defp get_async_pids(private) do
+    case private do
+      %{phoenix_async: ref_pids} -> Map.values(ref_pids)
+      %{} -> []
+    end
   end
 end
diff --git a/test/phoenix_live_view_test.exs b/test/phoenix_live_view_test.exs
index 38a772eea9..8e8a34600f 100644
--- a/test/phoenix_live_view_test.exs
+++ b/test/phoenix_live_view_test.exs
@@ -288,4 +288,17 @@ defmodule Phoenix.LiveViewUnitTest do
                {:live, :patch, %{kind: :push, to: "/counter/123"}}
     end
   end
+
+  describe "put_private" do
+    test "assigns private keys" do
+      assert @socket.private[:hello] == nil
+      assert put_private(@socket, :hello, "world").private[:hello] == "word"
+    end
+
+    test "disallows reserved keys" do
+      assert_raise ArgumentError, ~r/reserved/, fn ->
+        put_private(@socket, :assign_new, "boom")
+      end
+    end
+  end
 end
diff --git a/test/support/live_views/general.ex b/test/support/live_views/general.ex
index ab343edc4b..90a7b6b872 100644
--- a/test/support/live_views/general.ex
+++ b/test/support/live_views/general.ex
@@ -339,9 +339,9 @@ end
 
 defmodule Phoenix.LiveViewTest.AsyncLive do
   use Phoenix.LiveView
-  import Phoenix.LiveView.AsyncAssign
+  import Phoenix.LiveView.AsyncResult
 
-  on_mount {__MODULE__, :defaults}
+  on_mount({__MODULE__, :defaults})
 
   def on_mount(:defaults, _params, _session, socket) do
     {:cont, assign(socket, enum: false, lc: false)}
@@ -350,14 +350,15 @@ defmodule Phoenix.LiveViewTest.AsyncLive do
   def render(assigns) do
     ~H"""
     <.live_component :if={@lc} module={Phoenix.LiveViewTest.AsyncLive.LC} test={@lc} id="lc" />
-    <div :if={@async.data.loading?}>data loading...</div>
-    <div :if={@async.data.canceled?}>data canceled</div>
-    <div :if={!@async.data.loading? && @async.data.result == nil}>no data found</div>
-    <div :if={!@async.data.loading? && @async.data.result}>data: <%= inspect(@async.data.result) %></div>
-    <div :if={err = @async.data.error}>error: <%= inspect(err) %></div>
-
+    <div :if={@data.state == :loading}>data loading...</div>
+    <div :if={@data.state == :canceled}>data canceled</div>
+    <div :if={@data.state == :ok && @data.result == nil}>no data found</div>
+    <div :if={@data.state == :ok && @data.result}>data: <%= inspect(@data.result) %></div>
+    <%= with {kind, reason} when kind in [:error, :exit, :throw] <- @data.state do %>
+      <div><%= kind %>: <%= inspect(reason) %></div>
+    <% end %>
     <%= if @enum do %>
-      <div :for={i <- @async.data}><%= i %></div>
+      <div :for={i <- @data}><%= i %></div>
     <% end %>
     """
   end
@@ -412,7 +413,9 @@ defmodule Phoenix.LiveViewTest.AsyncLive do
 
   def handle_info(:boom, _socket), do: exit(:boom)
 
-  def handle_info(:cancel, socket), do: {:noreply, cancel_async(socket, :data)}
+  def handle_info(:cancel, socket) do
+    {:noreply, cancel_async(socket, socket.assigns.data)}
+  end
 
   def handle_info(:renew_canceled, socket) do
     {:noreply,
@@ -425,22 +428,22 @@ end
 
 defmodule Phoenix.LiveViewTest.AsyncLive.LC do
   use Phoenix.LiveComponent
-  import Phoenix.LiveView.AsyncAssign
+  import Phoenix.LiveView.AsyncResult
 
   def render(assigns) do
     ~H"""
     <div>
       <%= if @enum do %>
-        <div :for={i <- @async.lc_data}><%= i %></div>
+        <div :for={i <- @lc_data}><%= i %></div>
       <% end %>
-      <.async_result :let={data} assign={@async.lc_data}>
+      <AsyncResult.with_state :let={data} assign={@lc_data}>
         <:loading>lc_data loading...</:loading>
         <:canceled>lc_data canceled</:canceled>
         <:empty :let={_res}>no lc_data found</:empty>
-        <:error :let={err}>error: <%= inspect(err) %></:error>
+        <:error :let={{kind, reason}}><%= kind %>: <%= inspect(reason) %></:error>
 
         lc_data: <%= inspect(data) %>
-      </.async_result>
+      </AsyncResult.with_state>
     </div>
     """
   end
@@ -494,7 +497,9 @@ defmodule Phoenix.LiveViewTest.AsyncLive.LC do
 
   def update(%{action: :boom}, _socket), do: exit(:boom)
 
-  def update(%{action: :cancel}, socket), do: {:ok, cancel_async(socket, :lc_data)}
+  def update(%{action: :cancel}, socket) do
+    {:ok, cancel_async(socket, socket.assigns.lc_data)}
+  end
 
   def update(%{action: :renew_canceled}, socket) do
     {:ok,