ch2.erl

%%% ch2
%%%
%%% consistent hashing of static partitions of keys onto a dynamic
%%% list of members with weights
-module(ch2).

%% api
-export([partition_ring/1,
         select/2
        ]).
%% dummy testing
-export([start/0]).

%%-include_lib("proper/include/proper.hrl").
-include_lib("eunit/include/eunit.hrl").


-define(RINGTOP, trunc(math:pow(2,160)-1)).  % SHA-1 space
-define(PARTITIONS, 4096).


-type index()          :: integer().

-type member()         :: {member_id(), member_weight()}.
-type member_id()      :: term().
-type member_weight()  :: integer().

-type ring(X)          :: [{index(), X}].
-type partition_ring() :: ring(member_id()).
-type member_ring()    :: ring(member_id()).


%%% api ----------

-spec select(term(), partition_ring()) -> member_id().
select(Term, Ring) ->
  hd(ordered_from(to_index(Term), Ring)).

-spec partition_ring([member()]) -> partition_ring().
partition_ring(Members) ->
  MemberRing = ring_from_members(Members),
  Inc = ?RINGTOP div ?PARTITIONS,
  Partitions = [Index || Index <- lists:seq(0, (?RINGTOP - 1), Inc)],
  assign_members(Partitions, MemberRing, []).



%%% internal ----------

-spec assign_members([index()], member_ring(), partition_ring()) ->
                        partition_ring().
assign_members([], _, PartitionRing) ->
  lists:reverse(PartitionRing);
%% all partitions bigger than the last member wrap around to first member
assign_members(Partitions, [], PartitionRing0) ->
  PartitionRing = lists:reverse(PartitionRing0),
  {_, FirstMember} = hd(PartitionRing),
  PartitionRing ++ [{P, FirstMember} || P <- Partitions];
assign_members([P | Ps], [{I, M} | Ms], PartitionRing) ->
  case P =< I of
    true ->
      assign_members(Ps,       [{I, M} | Ms], [{P, M} | PartitionRing]);
    false ->
      assign_members([P | Ps], Ms,            PartitionRing)
  end.


-spec ring_from_members([member()]) -> member_ring().
ring_from_members(Members) ->
  lists:sort(
    lists:flatten(
      [[{to_index({MemberId, I}), MemberId} ||
         I <- lists:seq(1, MemberWeight * ?PARTITIONS)] ||
        {MemberId, MemberWeight} <- Members])).


-spec ordered_from(index(), ring(X)) -> [X].
ordered_from(Index, PartitionRing) ->
  Inc = ?RINGTOP div ?PARTITIONS,
  {A, B} = lists:split((Index div Inc)+1, PartitionRing),
  B ++ A.

-spec to_index(term()) -> index().
to_index(Term) ->
  <<IndexAsInt:160/integer>> = crypto:sha(term_to_binary(Term)),
  IndexAsInt.



%%% test

partitions_per_member(Ring) ->
  L = lists:sort([{M, P} || {P, M} <- Ring]),
  [{M, length(Ps) / ?PARTITIONS } || {M, Ps} <- group_partitions(L, [])].
  
group_partitions([], Acc) ->
  lists:reverse(Acc);
group_partitions([{M, P} | Rest], [{M, Ps} | Acc]) ->
  group_partitions(Rest, [{M, [P | Ps]} | Acc]);
group_partitions([{M, P} | Rest], Acc) ->
  group_partitions(Rest, [{M, [P]} | Acc]).


diff_rings(Ring1, Ring2) ->
  [{I, {M1, M2}} || {{I, M1}, {I, M2}} <- lists:zip(Ring1, Ring2), M1 =/= M2].


%% dummy testing

start() ->
  io:format("dummy test starting~n"),
  DummyMembers = dummy_members(3, 1),
  ?debugFmt("members: ~p~n", [DummyMembers]),
  %% MemberRing = ring_from_members(DummyMembers),
  %% ?debugFmt("ring: ~p~n", [MemberRing]),
  %% PartitionRing = partition_ring(DummyMembers),
  %% ?debugFmt("ring: ~p~n", [PartitionRing]),
  Ring1 = partition_ring(dummy_members(8, 1)),
  Ring2 = partition_ring(dummy_members(7, 1) ++ [{member8, 2}]),
  ?debugFmt("parts per member ring1: ~p~n", [partitions_per_member(Ring1)]),
  ?debugFmt("parts per member ring2: ~p~n", [partitions_per_member(Ring2)]),
  RingDiff = length(diff_rings(Ring1, Ring2)) / ?PARTITIONS,
  ?debugFmt("ringdiff: ~f~n", [RingDiff]),
  Member = select("test", Ring1),
  Member1 = select("test1", Ring1),
  ?debugFmt("index_of(\"test\": ~p~n", [{to_index("test"), to_index("test1")}]),
  ?debugFmt("select(\"test\": ~p~n", [{Member, Member1}]),
  io:format("dummy test done~n").

dummy_members(N, W) ->
  [{list_to_atom("member" ++ integer_to_list(M)), W} || M <- lists:seq(1, N)].