From afb5eb32fddb1984cec669c9da2bc792e8b584b1 Mon Sep 17 00:00:00 2001
From: Tony Winn <hi@tonywinn.me>
Date: Fri, 13 Oct 2023 01:41:07 -0700
Subject: [PATCH] Partitioned adapter supports two item tuples as keys (#214)

There was a problem that `get_all` when asking for a two element tuple key (ie
`{"foo", 1}) would only consider the first element of the tuple to be the key,
and therefore may ask for the data on the wrong node. This seems to have been
done to support put_(new_)all which puts the key in the first element and values in
the second.

I've passed along the action to the `group_keys_by_node` function so we're able
to know which structure to expect.
---
 lib/nebulex/adapters/partitioned.ex        | 18 ++++++++++++------
 test/nebulex/adapters/partitioned_test.exs | 22 ++++++++++++++++++++++
 2 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/lib/nebulex/adapters/partitioned.ex b/lib/nebulex/adapters/partitioned.ex
index 173acde7..d053fd85 100644
--- a/lib/nebulex/adapters/partitioned.ex
+++ b/lib/nebulex/adapters/partitioned.ex
@@ -711,18 +711,24 @@ defmodule Nebulex.Adapters.Partitioned do
     end
   end
 
-  defp group_keys_by_node(enum, adapter_meta) do
+  defp group_keys_by_node(enum, adapter_meta, :get_all) do
     Enum.reduce(enum, %{}, fn
-      {key, _} = entry, acc ->
-        node = get_node(adapter_meta, key)
-        Map.put(acc, node, [entry | Map.get(acc, node, [])])
-
       key, acc ->
         node = get_node(adapter_meta, key)
         Map.put(acc, node, [key | Map.get(acc, node, [])])
     end)
   end
 
+  @put_all_actions [:put_all, :put_new_all]
+  defp group_keys_by_node(enum, adapter_meta, put_all_action)
+       when put_all_action in @put_all_actions do
+    Enum.reduce(enum, %{}, fn
+      {key, _} = entry, acc ->
+        node = get_node(adapter_meta, key)
+        Map.put(acc, node, [entry | Map.get(acc, node, [])])
+    end)
+  end
+
   defp map_reduce(
          enum,
          %{task_sup: task_sup} = meta,
@@ -733,7 +739,7 @@ defmodule Nebulex.Adapters.Partitioned do
        ) do
     groups =
       enum
-      |> group_keys_by_node(meta)
+      |> group_keys_by_node(meta, action)
       |> Enum.map(fn {node, group} ->
         {node, {__MODULE__, :with_dynamic_cache, [meta, action, [group | args]]}}
       end)
diff --git a/test/nebulex/adapters/partitioned_test.exs b/test/nebulex/adapters/partitioned_test.exs
index dc0bee95..abf156c7 100644
--- a/test/nebulex/adapters/partitioned_test.exs
+++ b/test/nebulex/adapters/partitioned_test.exs
@@ -110,6 +110,28 @@ defmodule Nebulex.Adapters.PartitionedTest do
       end)
     end
 
+    test "custom keyslot supports two item tuple keys for get_all" do
+      defmodule TupleKeyslot do
+        @behaviour Nebulex.Adapter.Keyslot
+
+        @impl true
+        def hash_slot({_, _} = key, range) do
+          key
+          |> :erlang.phash2()
+          |> rem(range)
+        end
+      end
+
+      test_with_dynamic_cache(
+        Partitioned,
+        [name: :custom_keyslot_with_tuple_keys, keyslot: TupleKeyslot],
+        fn ->
+          assert Partitioned.put_all([{{"foo", 1}, "bar"}]) == :ok
+          assert Partitioned.get_all([{"foo", 1}]) == %{{"foo", 1} => "bar"}
+        end
+      )
+    end
+
     test "get_and_update" do
       assert Partitioned.get_and_update(1, &Partitioned.get_and_update_fun/1) == {nil, 1}
       assert Partitioned.get_and_update(1, &Partitioned.get_and_update_fun/1) == {1, 2}