change structure of Polyjuice.Client process tree

4 files changed, 134 insertions, 127 deletions

diff --git a/lib/polyjuice/client.ex b/lib/polyjuice/client.ex
index 3f1668c..a39ac08 100644
--- a/lib/polyjuice/client.ex
+++ b/lib/polyjuice/client.ex
@@ -28,9 +28,10 @@ defmodule Polyjuice.Client do
   The client defined in this module should work for most cases.  If you want
   more control, you can use `Polyjuice.Client.LowLevel` instead.
 
-  To start a client with this module, create a process using `start_link/2`,
-  and then call `get_client/1` to get a struct that can be used with the above
-  modules.  To stop the client, use `Polyjuice.Client.API.stop/3`.
+  To start a client with this module, create a process using `start/2` or
+  `start_link/2`, and then call `get_client/1` to get a struct that can be used
+  with the above modules.  To stop the client, use
+  `Polyjuice.Client.API.stop/3`.
 
   """
 
@@ -44,13 +45,15 @@ defmodule Polyjuice.Client do
       Polyjuice.Client.child_spec(["http://localhost:8008", sync: false])
 
   """
-  use Supervisor
+  use GenServer
   require Logger
 
   @typedoc """
   Matrix client.
 
-  This struct is created by `start_link/2`.
+  This struct can be obtained by calling `get_client/1` after the client
+  process has been started with `start/2` or `start_link/2`, and implements
+  `Polyjuice.Client.API` protocol.
   """
   # - `base_url`: Required.  The base URL for the homeserver.
   # - `id`: An ID for the client.
@@ -60,6 +63,7 @@ defmodule Polyjuice.Client do
   @opaque t :: %__MODULE__{
             base_url: String.t(),
             id: integer,
+            pid: pid,
             storage: Polyjuice.Client.Storage.t(),
             handler: Polyjuice.Client.Handler.t(),
             opts: list,
@@ -71,6 +75,7 @@ defmodule Polyjuice.Client do
   defstruct [
     :base_url,
     :id,
+    :pid,
     :storage,
     :handler,
     sync: true,
@@ -95,28 +100,28 @@ defmodule Polyjuice.Client do
   - `sync_filter`: the filter to use for the sync.  Defaults to no filter.
   """
   @spec start_link(base_url :: String.t(), opts :: Keyword.t()) :: {:ok, pid}
-  def start_link(base_url, opts) when is_binary(base_url) and is_list(opts) do
-    base_url =
-      if(String.ends_with?(base_url, "/"), do: base_url, else: base_url <> "/")
-      |> URI.parse()
+  def start_link(base_url, opts \\ []) when is_binary(base_url) and is_list(opts) do
+    GenServer.start_link(__MODULE__, [base_url, opts])
+  end
 
-    client_id = Agent.get_and_update(Polyjuice.Client.ID, &{&1, &1 + 1})
+  @doc """
+  Start a client.
 
-    Supervisor.start_link(
-      __MODULE__,
-      [client_id, base_url, opts],
-      name: process_name(client_id, :supervisor)
-    )
+  See `start_link/2`.
+  """
+  @spec start(base_url :: String.t(), opts :: Keyword.t()) :: {:ok, pid}
+  def start(base_url, opts \\ []) when is_binary(base_url) and is_list(opts) do
+    GenServer.start(__MODULE__, [base_url, opts])
   end
 
-  def start_link(base_url) when is_binary(base_url), do: start_link(base_url, [])
+  @impl GenServer
+  def init([base_url, opts]) do
+    base_url =
+      if(String.ends_with?(base_url, "/"), do: base_url, else: base_url <> "/")
+      |> URI.parse()
 
-  def start_link([base_url | opts]) when is_binary(base_url) and is_list(opts) do
-    start_link(base_url, opts)
-  end
+    client_id = Agent.get_and_update(Polyjuice.Client.ID, &{&1, &1 + 1})
 
-  @impl Supervisor
-  def init([client_id, base_url, opts]) do
     sync = Keyword.get(opts, :sync, true)
     storage = Keyword.get(opts, :storage)
     handler = Keyword.get(opts, :handler)
@@ -171,50 +176,81 @@ defmodule Polyjuice.Client do
           x
       end
 
-    children = [
-      %{
-        id: Polyjuice.Client,
-        start:
-          {Agent, :start_link,
-           [
-             fn ->
-               %{
-                 access_token: access_token,
-                 user_id: user_id,
-                 device_id: device_id,
-                 client: %__MODULE__{
-                   base_url: base_url,
-                   id: client_id,
-                   storage: storage,
-                   handler: handler,
-                   sync: sync,
-                   opts: opts,
-                   test: Keyword.get(opts, :test, false)
-                 }
-               }
-             end,
-             [name: process_name(client_id, :state)]
-           ]}
-      }
-      | if(sync and access_token != nil and handler != nil and storage != nil,
-          do: [sync_child_spec(base_url, client_id, opts)],
-          else: []
+    {:ok, supervisor} =
+      DynamicSupervisor.start_link(
+        strategy: :one_for_one,
+        name: process_name(client_id, :supervisor)
+      )
+
+    if sync and access_token != nil and handler != nil and storage != nil do
+      {:ok, _pid} =
+        DynamicSupervisor.start_child(
+          supervisor,
+          sync_child_spec(base_url, client_id, self(), opts)
         )
-    ]
+    end
 
-    Supervisor.init(children, strategy: :rest_for_one)
+    {:ok,
+     %{
+       access_token: access_token,
+       user_id: user_id,
+       device_id: device_id,
+       supervisor: supervisor,
+       client: %__MODULE__{
+         base_url: base_url,
+         id: client_id,
+         pid: self(),
+         storage: storage,
+         handler: handler,
+         sync: sync,
+         opts: opts,
+         test: Keyword.get(opts, :test, false)
+       }
+     }}
+  end
+
+  @doc """
+  Returns a specification to start this under a supervisor.
+
+  `arg` is a list where the first element is the server's base URL, and the
+  remainder of the list are options as document in `start_link/2`.
+  """
+  def child_spec([base_url | opts]) do
+    %{
+      id: __MODULE__,
+      start: {__MODULE__, :start_link, [base_url, opts]},
+      restart: :transient
+    }
   end
 
   @doc """
   Get a struct that implements `Polyjuice.Client.API` from the pid given by
-  `start_link`.
+  `start/2` or `start_link/2`.
   """
   def get_client(pid) do
-    agent_pid =
-      Supervisor.which_children(pid)
-      |> Enum.find_value(fn {id, pid, _, _} -> if id == Polyjuice.Client, do: pid end)
+    GenServer.call(pid, :get_client)
+  end
+
+  @impl GenServer
+  def terminate(_reason, %{supervisor: supervisor}) do
+    DynamicSupervisor.stop(supervisor)
+  end
+
+  @impl GenServer
+  def handle_call(:get_client, _from, %{client: client} = state) do
+    {:reply, client, state}
+  end
+
+  @impl GenServer
+  def handle_call(:get_state, _from, state) do
+    {:reply, Map.take(state, [:access_token, :user_id, :device_id]), state}
+  end
 
-    Agent.get(agent_pid, &Map.fetch!(&1, :client))
+  @impl GenServer
+  def handle_cast({:set, new_state}, state) do
+    Map.take(new_state, [:access_token, :user_id, :device_id])
+    |> (&Map.merge(state, &1)).()
+    |> (&{:noreply, &1}).()
   end
 
   @doc false
@@ -222,11 +258,11 @@ defmodule Polyjuice.Client do
     {:via, Registry, {Polyjuice.Client, {id, process}}}
   end
 
-  defp sync_child_spec(base_url, client_id, opts) do
+  defp sync_child_spec(base_url, client_id, pid, opts) do
     %{
       id: Polyjuice.Client.Sync,
       restart: :transient,
-      start: {Task, :start_link, [Polyjuice.Client.Sync, :sync, [base_url, client_id, opts]]}
+      start: {Task, :start_link, [Polyjuice.Client.Sync, :sync, [base_url, client_id, pid, opts]]}
     }
   end
 
@@ -281,7 +317,7 @@ defmodule Polyjuice.Client do
 
   defimpl Polyjuice.Client.API do
     def call(
-          %{base_url: base_url, id: id, test: test} = client,
+          %{base_url: base_url, id: id, pid: pid, test: test} = client,
           endpoint
         ) do
       %Polyjuice.Client.Endpoint.HttpSpec{
@@ -297,7 +333,7 @@ defmodule Polyjuice.Client do
 
       access_token =
         if auth_required do
-          Agent.get(Polyjuice.Client.process_name(id, :state), &Map.get(&1, :access_token))
+          GenServer.call(pid, :get_state) |> Map.get(:access_token)
         else
           nil
         end
@@ -351,7 +387,7 @@ defmodule Polyjuice.Client do
                     Polyjuice.Client.kill_sync(id)
 
                     Polyjuice.Client.set_logged_out(
-                      id,
+                      pid,
                       client.storage,
                       client.handler,
                       Map.get(json, "soft_logout", false)
@@ -403,8 +439,8 @@ defmodule Polyjuice.Client do
       "#{Node.self()}_#{:erlang.system_time(:millisecond)}_#{:erlang.unique_integer()}"
     end
 
-    def stop(%{id: id}, reason \\ :normal, timeout \\ :infinity) do
-      Supervisor.stop(Polyjuice.Client.process_name(id, :supervisor), reason, timeout)
+    def stop(%{pid: pid}, reason \\ :normal, timeout \\ :infinity) do
+      GenServer.stop(pid, reason, timeout)
     end
   end
 
@@ -519,9 +555,9 @@ defmodule Polyjuice.Client do
           {:ok,
            %{"access_token" => access_token, "user_id" => user_id, "device_id" => device_id} =
                http_ret} ->
-        Agent.cast(
-          process_name(client.id, :state),
-          &%{&1 | access_token: access_token, user_id: user_id, device_id: device_id}
+        GenServer.cast(
+          client.pid,
+          {:set, %{access_token: access_token, user_id: user_id, device_id: device_id}}
         )
 
         if client.storage do
@@ -561,9 +597,9 @@ defmodule Polyjuice.Client do
 
           supervisor_name = process_name(client.id, :supervisor)
 
-          Supervisor.start_child(
+          DynamicSupervisor.start_child(
             supervisor_name,
-            sync_child_spec(client.base_url, client.id, client.opts)
+            sync_child_spec(client.base_url, client.id, client.pid, client.opts)
           )
         end
 
@@ -578,7 +614,7 @@ defmodule Polyjuice.Client do
   Log out an existing session.
   """
   @spec log_out(client :: Polyjuice.Client.t()) :: {:ok} | any
-  def log_out(%Polyjuice.Client{id: id, storage: storage, handler: handler} = client) do
+  def log_out(%Polyjuice.Client{id: id, pid: pid, storage: storage, handler: handler} = client) do
     kill_sync(id)
 
     case Polyjuice.Client.API.call(
@@ -586,7 +622,7 @@ defmodule Polyjuice.Client do
            %Polyjuice.Client.Endpoint.PostLogout{}
          ) do
       {:ok} ->
-        set_logged_out(id, storage, handler, false)
+        set_logged_out(pid, storage, handler, false)
 
         if storage do
           Polyjuice.Client.Storage.kv_del(storage, "ca.uhoreg.polyjuice", "user_id")
@@ -603,13 +639,19 @@ defmodule Polyjuice.Client do
   @doc false
   def kill_sync(id) do
     supervisor_name = process_name(id, :supervisor)
-    Supervisor.terminate_child(supervisor_name, Polyjuice.Client.Sync)
-    Supervisor.delete_child(supervisor_name, Polyjuice.Client.Sync)
+
+    case DynamicSupervisor.which_children(supervisor_name) do
+      [{_, pid, _, _}] when is_pid(pid) ->
+        DynamicSupervisor.terminate_child(supervisor_name, pid)
+
+      _ ->
+        nil
+    end
   end
 
   @doc false
-  def set_logged_out(id, storage, handler, soft_logout) do
-    Agent.cast(process_name(id, :state), &%{&1 | access_token: nil})
+  def set_logged_out(pid, storage, handler, soft_logout) do
+    GenServer.cast(pid, {:set, %{access_token: nil}})
 
     if storage do
       Polyjuice.Client.Storage.kv_del(storage, "ca.uhoreg.polyjuice", "access_token")
diff --git a/lib/polyjuice/client/sync.ex b/lib/polyjuice/client/sync.ex
index 3ac6f8a..573af71 100644
--- a/lib/polyjuice/client/sync.ex
+++ b/lib/polyjuice/client/sync.ex
@@ -26,6 +26,7 @@ defmodule Polyjuice.Client.Sync do
     :handler,
     :conn_ref,
     :id,
+    :pid,
     :homeserver_url,
     :uri,
     :storage,
@@ -45,6 +46,7 @@ defmodule Polyjuice.Client.Sync do
   def sync(
         homeserver_url,
         id,
+        pid,
         opts
       ) do
     storage = Keyword.fetch!(opts, :storage)
@@ -92,6 +94,7 @@ defmodule Polyjuice.Client.Sync do
     connect(%__MODULE__{
       handler: handler,
       id: id,
+      pid: pid,
       homeserver_url: homeserver_url,
       uri: uri,
       query_params: query_params,
@@ -138,13 +141,7 @@ defmodule Polyjuice.Client.Sync do
 
     e = &URI.encode_www_form/1
 
-    {access_token, user_id} =
-      Agent.get(
-        Polyjuice.Client.process_name(state.id, :state),
-        fn %{access_token: access_token, user_id: user_id} ->
-          {access_token, user_id}
-        end
-      )
+    %{access_token: access_token, user_id: user_id} = GenServer.call(state.pid, :get_state)
 
     path =
       URI.merge(
@@ -202,7 +199,7 @@ defmodule Polyjuice.Client.Sync do
               :hackney.close(state.conn_ref)
 
               Polyjuice.Client.set_logged_out(
-                state.id,
+                state.pid,
                 state.storage,
                 state.handler,
                 Map.get(json_body, "soft_logout", false)
@@ -245,11 +242,7 @@ defmodule Polyjuice.Client.Sync do
   defp do_sync(state) do
     if state.backoff, do: :timer.sleep(state.backoff * 1000)
 
-    access_token =
-      Agent.get(
-        Polyjuice.Client.process_name(state.id, :state),
-        fn %{access_token: access_token} -> access_token end
-      )
+    %{access_token: access_token} = GenServer.call(state.pid, :get_state)
 
     headers = [
       {"Accept", "application/json"},
@@ -298,7 +291,7 @@ defmodule Polyjuice.Client.Sync do
               :hackney.close(state.conn_ref)
 
               Polyjuice.Client.set_logged_out(
-                state.id,
+                state.pid,
                 state.storage,
                 state.handler,
                 Map.get(json_body, "soft_logout", false)
@@ -429,7 +422,7 @@ defmodule Polyjuice.Client.Sync do
         end
       )
 
-    user_id = Agent.get(Polyjuice.Client.process_name(state.id, :state), &Map.get(&1, :user_id))
+    %{user_id: user_id} = GenServer.call(state.pid, :get_state)
 
     inviter =
       invite_state
diff --git a/test/polyjuice/client/sync_test.exs b/test/polyjuice/client/sync_test.exs
index e1085d6..94387aa 100644
--- a/test/polyjuice/client/sync_test.exs
+++ b/test/polyjuice/client/sync_test.exs
@@ -440,9 +440,7 @@ defmodule Polyjuice.Client.SyncTest do
       Process.sleep(10)
 
       assert Polyjuice.Client.process_name(client.id, :supervisor)
-             |> Supervisor.which_children()
-             |> Enum.find_value(fn {id, pid, _, _} -> if id == Polyjuice.Client.Sync, do: pid end) ==
-               :undefined
+             |> DynamicSupervisor.which_children() == []
 
       Polyjuice.Client.API.stop(client)
 
@@ -495,9 +493,7 @@ defmodule Polyjuice.Client.SyncTest do
       Process.sleep(10)
 
       assert Polyjuice.Client.process_name(client.id, :supervisor)
-             |> Supervisor.which_children()
-             |> Enum.find_value(fn {id, pid, _, _} -> if id == Polyjuice.Client.Sync, do: pid end) ==
-               :undefined
+             |> DynamicSupervisor.which_children() == []
 
       Polyjuice.Client.API.stop(client)
 
diff --git a/test/polyjuice/client_test.exs b/test/polyjuice/client_test.exs
index d9f95a8..96dc829 100644
--- a/test/polyjuice/client_test.exs
+++ b/test/polyjuice/client_test.exs
@@ -235,10 +235,7 @@ defmodule Polyjuice.ClientTest do
 
       Polyjuice.Client.log_in_with_password(client, "@alice:example.org", "password")
 
-      assert Agent.get(
-               Polyjuice.Client.process_name(client.id, :state),
-               &Map.take(&1, [:access_token, :user_id, :device_id])
-             ) == %{
+      assert GenServer.call(client.pid, :get_state) == %{
                access_token: "m.id.user_login",
                user_id: "@alice:example.org",
                device_id: "foo"
@@ -264,10 +261,7 @@ defmodule Polyjuice.ClientTest do
       assert Polyjuice.Client.Storage.kv_get(storage, "ca.uhoreg.polyjuice", "user_id") == nil
       assert Polyjuice.Client.Storage.kv_get(storage, "ca.uhoreg.polyjuice", "device_id") == nil
 
-      assert Agent.get(
-               Polyjuice.Client.process_name(client.id, :state),
-               &Map.fetch!(&1, :access_token)
-             ) == nil
+      assert GenServer.call(client.pid, :get_state) |> Map.get(:access_token) == nil
 
       Polyjuice.Client.log_in_with_password(
         client,
@@ -277,10 +271,7 @@ defmodule Polyjuice.ClientTest do
 
       assert_receive({:polyjuice_client, :logged_in, {"@alice:example.org", "foo", _}})
 
-      assert Agent.get(
-               Polyjuice.Client.process_name(client.id, :state),
-               &Map.take(&1, [:access_token, :user_id, :device_id])
-             ) == %{
+      assert GenServer.call(client.pid, :get_state) == %{
                access_token: "m.id.thirdparty_login",
                user_id: "@alice:example.org",
                device_id: "foo"
@@ -294,10 +285,7 @@ defmodule Polyjuice.ClientTest do
 
       assert_receive({:polyjuice_client, :logged_in, {"@alice:example.org", "foo", _}})
 
-      assert Agent.get(
-               Polyjuice.Client.process_name(client.id, :state),
-               &Map.take(&1, [:access_token, :user_id, :device_id])
-             ) == %{
+      assert GenServer.call(client.pid, :get_state) == %{
                access_token: "m.id.phone_login",
                user_id: "@alice:example.org",
                device_id: "foo"
@@ -313,10 +301,7 @@ defmodule Polyjuice.ClientTest do
 
       assert_receive({:polyjuice_client, :logged_in, {"@alice:example.org", "foo", _}})
 
-      assert Agent.get(
-               Polyjuice.Client.process_name(client.id, :state),
-               &Map.take(&1, [:access_token, :user_id, :device_id])
-             ) == %{
+      assert GenServer.call(client.pid, :get_state) == %{
                access_token: "ca.uhoreg.foo_login",
                user_id: "@alice:example.org",
                device_id: "foo"
@@ -349,10 +334,7 @@ defmodule Polyjuice.ClientTest do
 
     client = Polyjuice.Client.get_client(client_pid)
 
-    assert Agent.get(
-             Polyjuice.Client.process_name(client.id, :state),
-             &Map.take(&1, [:access_token, :user_id, :device_id])
-           ) == %{
+    assert GenServer.call(client.pid, :get_state) == %{
              access_token: "an_access_token",
              user_id: "@alice:example.org",
              device_id: "a_device"
@@ -382,10 +364,7 @@ defmodule Polyjuice.ClientTest do
 
     client2 = Polyjuice.Client.get_client(client2_pid)
 
-    assert Agent.get(
-             Polyjuice.Client.process_name(client2.id, :state),
-             &Map.take(&1, [:access_token, :user_id, :device_id])
-           ) == %{
+    assert GenServer.call(client2.pid, :get_state) == %{
              access_token: "an_access_token",
              user_id: "@alice:example.org",
              device_id: "a_device"
@@ -431,10 +410,7 @@ defmodule Polyjuice.ClientTest do
 
       Polyjuice.Client.Room.send_event(client, "!a_room:example.org", "m.room.message", %{})
 
-      assert Agent.get(
-               Polyjuice.Client.process_name(client.id, :state),
-               &Map.fetch!(&1, :access_token)
-             ) == nil
+      assert GenServer.call(client.pid, :get_state) |> Map.get(:state) == nil
 
       assert Polyjuice.Client.Storage.kv_get(storage, "ca.uhoreg.polyjuice", "access_token") ==
                nil