Repository: Kraigie/nostrum
Pull request #486 (2023-05-09 22:51:52)

Add `get_with_user/2` to the MemberCache


Patch:
diff --git a/lib/nostrum/cache/member_cache.ex b/lib/nostrum/cache/member_cache.ex
index 99dce640..73a264ce 100644
--- a/lib/nostrum/cache/member_cache.ex
+++ b/lib/nostrum/cache/member_cache.ex
@@ -112,6 +112,26 @@ defmodule Nostrum.Cache.MemberCache do
   defdelegate get(guild_id, member_id), to: @configured_cache
   defdelegate by_user(member_id), to: @configured_cache
 
+  @doc """
+  Return a member together with its user via the user cache.
+  """
+  @spec get_with_user(Guild.id(), Member.user_id()) :: {Member.t(), User.t() | nil} | nil
+  def get_with_user(guild_id, member_id) do
+    case get(guild_id, member_id) do
+      {:ok, member} ->
+        case UserCache.get(member_id) do
+          {:ok, user} ->
+            {member, user}
+
+          _ ->
+            {member, nil}
+        end
+
+      _ ->
+        nil
+    end
+  end
+
   @doc """
   Return an enumerable of members with their users.
 
@@ -130,13 +150,6 @@ defmodule Nostrum.Cache.MemberCache do
   """
   @spec get_with_users(Guild.id()) :: Enumerable.t({Member.t(), User.t()})
   def get_with_users(guild_id) do
-    member_query_handle = qlc_handle()
-    user_query_handle = UserCache.qlc_handle()
-    initial_bindings = :erl_eval.new_bindings()
-    member_bindings = :erl_eval.add_binding(:MemberHandle, member_query_handle, initial_bindings)
-    user_bindings = :erl_eval.add_binding(:UserHandle, user_query_handle, member_bindings)
-    final_bindings = :erl_eval.add_binding(:RequestedGuildId, guild_id, user_bindings)
-
     joined_handle =
       :qlc.string_to_handle(
         '[{Member, User} || ' ++
@@ -145,7 +158,9 @@ defmodule Nostrum.Cache.MemberCache do
           'GuildId =:= RequestedGuildId, ' ++
           'UserId =:= MemberId].',
         [],
-        final_bindings
+        MemberHandle: qlc_handle(),
+        UserHandle: UserCache.qlc_handle(),
+        RequestedGuildId: guild_id
       )
 
     Stream.resource(
diff --git a/lib/nostrum/shard/dispatch.ex b/lib/nostrum/shard/dispatch.ex
index 1ac22c40..7dacfb20 100644
--- a/lib/nostrum/shard/dispatch.ex
+++ b/lib/nostrum/shard/dispatch.ex
@@ -362,10 +362,6 @@ defmodule Nostrum.Shard.Dispatch do
           voice.channel_id != p.channel_id and is_pid(voice.session_pid) ->
             Voice.on_channel_join_change(p, voice)
 
-          # Already in this channel but connection died:
-          is_pid(voice.session_pid) and not Process.alive?(voice.session_pid) ->
-            Voice.restart_session(p)
-
           # Already in this channel:
           true ->
             :noop
diff --git a/lib/nostrum/voice.ex b/lib/nostrum/voice.ex
index 979eebd9..3737d99b 100644
--- a/lib/nostrum/voice.ex
+++ b/lib/nostrum/voice.ex
@@ -117,8 +117,8 @@ defmodule Nostrum.Voice do
   end
 
   @doc false
-  def remove_voice(guild_id, pre_cleanup_args \\ []) do
-    GenServer.cast(VoiceStateMap, {:remove, guild_id, pre_cleanup_args})
+  def remove_voice(guild_id) do
+    GenServer.cast(VoiceStateMap, {:remove, guild_id})
   end
 
   @doc """
@@ -748,7 +748,7 @@ defmodule Nostrum.Voice do
     voice =
       state
       |> Map.get(guild_id, VoiceState.new(guild_id: guild_id))
-      |> Map.merge(Map.new(args))
+      |> Map.merge(Enum.into(args, %{}))
 
     state = Map.put(state, guild_id, voice)
 
@@ -764,12 +764,8 @@ defmodule Nostrum.Voice do
   end
 
   @doc false
-  def handle_cast({:remove, guild_id, pre_cleanup_args}, state) do
-    state
-    |> Map.get(guild_id, %{})
-    |> Map.merge(Map.new(pre_cleanup_args))
-    |> VoiceState.cleanup()
-
+  def handle_cast({:remove, guild_id}, state) do
+    VoiceState.cleanup(state[guild_id])
     {:noreply, Map.delete(state, guild_id)}
   end
 
@@ -792,41 +788,51 @@ defmodule Nostrum.Voice do
 
   @doc false
   def on_channel_join_change(p, voice) do
+    v_ws = Session.get_ws_state(voice.session_pid)
+
+    # On the off-chance that we receive Voice Server Update first:
+    {new_token, new_gateway} =
+      if voice.token == v_ws.token do
+        # Need to reset
+        {nil, nil}
+      else
+        # Already updated
+        {voice.token, voice.gateway}
+      end
+
+    %{
+      ffmpeg_proc: ffmpeg_proc,
+      raw_audio: raw_audio,
+      raw_stateful: raw_stateful,
+      current_url: current_url,
+      persist_source: persist_source,
+      persist_playback: persist_playback
+    } = voice
+
     # Nil-ify ffmpeg_proc so it doesn't get closed when cleanup is called
-    if voice.persist_source do
-      remove_voice(p.guild_id, ffmpeg_proc: nil)
-    else
-      remove_voice(p.guild_id)
-    end
+    if persist_source, do: update_voice(p.guild_id, ffmpeg_proc: nil)
 
-    update_voice(
-      p.guild_id,
+    remove_voice(p.guild_id)
+
+    fields =
       [
         channel_id: p.channel_id,
         session: p.session_id,
         self_mute: p.self_mute,
         self_deaf: p.self_deaf,
-        token: nil,
-        gateway: nil
-      ] ++ persistent_args(voice)
-    )
-  end
-
-  @doc false
-  def restart_session(p) do
-    leave_channel(p.guild_id)
-    join_channel(p.guild_id, p.channel_id, p.self_mute, p.self_deaf)
-  end
-
-  defp persistent_args(%{persist_source: true} = voice) do
-    [
-      ffmpeg_proc: voice.ffmpeg_proc,
-      raw_audio: voice.raw_audio,
-      raw_stateful: voice.raw_stateful,
-      current_url: voice.current_url,
-      persist_playback: voice.persist_playback
-    ]
+        token: new_token,
+        gateway: new_gateway
+      ] ++
+        if persist_source,
+          do: [
+            ffmpeg_proc: ffmpeg_proc,
+            raw_audio: raw_audio,
+            raw_stateful: raw_stateful,
+            current_url: current_url,
+            persist_playback: persist_playback
+          ],
+          else: []
+
+    update_voice(p.guild_id, fields)
   end
-
-  defp persistent_args(_voice), do: []
 end
diff --git a/lib/nostrum/voice/session.ex b/lib/nostrum/voice/session.ex
index e3bcb872..122a3189 100644
--- a/lib/nostrum/voice/session.ex
+++ b/lib/nostrum/voice/session.ex
@@ -96,11 +96,6 @@ defmodule Nostrum.Voice.Session do
     end
   end
 
-  def handle_info({:gun_ws, _conn, _stream, :close}, state) do
-    Logger.info("Voice websocket closed (unknown reason)")
-    {:noreply, state}
-  end
-
   def handle_info({:gun_ws, _conn, _stream, {:close, errno, reason}}, state) do
     Logger.info("Voice websocket closed (errno #{errno}, reason #{inspect(reason)})")
 
@@ -215,10 +210,13 @@ defmodule Nostrum.Voice.Session do
     spawn(fn ->
       Process.monitor(state.conn_pid)
 
-      %VoiceState{} = voice = Voice.get_voice(state.guild_id)
+      %VoiceState{channel_id: chan, self_mute: mute, self_deaf: deaf} =
+        Voice.get_voice(state.guild_id)
 
       receive do
-        _ -> Voice.restart_session(voice)
+        _ ->
+          Voice.leave_channel(state.guild_id)
+          Voice.join_channel(state.guild_id, chan, mute, deaf)
       end
     end)
   end


Test patch:
diff --git a/test/nostrum/cache/member_cache_test.exs b/test/nostrum/cache/member_cache_test.exs
index 16bb70fc..31588df0 100644
--- a/test/nostrum/cache/member_cache_test.exs
+++ b/test/nostrum/cache/member_cache_test.exs
@@ -11,6 +11,40 @@ defmodule Nostrum.Cache.MemberCacheTest do
     Nostrum.Cache.MemberCache.ETS
   ]
 
+  describe "get_with_user/1" do
+    setup do
+      [member_pid: start_supervised!(MemberCache), user_pid: start_supervised!(UserCache)]
+    end
+
+    test "no member and no user" do
+      refute MemberCache.get_with_user(1234, 12059)
+    end
+
+    test "no member but user" do
+      user = %{id: :erlang.unique_integer()}
+      UserCache.create(user)
+      refute MemberCache.get_with_user(1234, user.id)
+    end
+
+    test "member but no user" do
+      guild_id = :erlang.unique_integer()
+      raw_member = %{user: %{id: :erlang.unique_integer()}}
+      assert created_member = MemberCache.create(guild_id, raw_member)
+      assert {^created_member, nil} = MemberCache.get_with_user(guild_id, raw_member.user.id)
+    end
+
+    test "member with user" do
+      guild_id = :erlang.unique_integer()
+      raw_user = %{id: :erlang.unique_integer()}
+      raw_member = %{user: raw_user}
+      assert created_user = UserCache.create(raw_user)
+      assert created_member = MemberCache.create(guild_id, raw_member)
+
+      assert {^created_member, ^created_user} =
+               MemberCache.get_with_user(guild_id, raw_member.user.id)
+    end
+  end
+
   describe "get_with_users/1" do
     setup do
       [member_pid: start_supervised!(MemberCache), user_pid: start_supervised!(UserCache)]