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}