|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| -module(seb_partition_mgr).
|
| -behaviour(gen_server).
|
|
|
| -export([start_link/1]).
|
| -export([init/1, handle_call/3, handle_cast/2, handle_info/2, terminate/2, code_change/3]).
|
|
|
|
|
| -export([assign_partition/2, get_partition_load/1, rebalance_partitions/0]).
|
|
|
| -define(SERVER, ?MODULE).
|
| -define(DEFAULT_PARTITIONS, 1024).
|
| -define(PARTITION_LOAD_THRESHOLD, 0.8).
|
|
|
| -record(state, {
|
| partition_count :: non_neg_integer(),
|
| partition_assignments :: map(),
|
| partition_loads :: map()
|
| }).
|
|
|
|
|
|
|
|
|
|
|
|
|
| -spec start_link(non_neg_integer()) -> gen_server:start_ret().
|
| start_link(PartitionCount) when is_integer(PartitionCount), PartitionCount > 0 ->
|
| gen_server:start_link({local, ?SERVER}, ?MODULE, [PartitionCount], []).
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| -spec assign_partition(binary(), atom()) -> non_neg_integer().
|
| assign_partition(AgentId, Competency) when is_binary(AgentId), is_atom(Competency) ->
|
| gen_server:call(?SERVER, {assign_partition, AgentId, Competency}).
|
|
|
|
|
|
|
|
|
|
|
| -spec get_partition_load(non_neg_integer()) -> float() | {error, not_found}.
|
| get_partition_load(Partition) when is_integer(Partition), Partition >= 0 ->
|
| gen_server:call(?SERVER, {get_partition_load, Partition}).
|
|
|
|
|
|
|
|
|
|
|
|
|
| -spec rebalance_partitions() -> ok | {error, term()}.
|
| rebalance_partitions() ->
|
| gen_server:call(?SERVER, rebalance_partitions).
|
|
|
|
|
|
|
|
|
|
|
|
|
| -spec init([non_neg_integer()]) -> {ok, #state{}}.
|
| init([PartitionCount]) ->
|
| State = #state{
|
| partition_count = PartitionCount,
|
| partition_assignments = maps:new(),
|
| partition_loads = init_loads(PartitionCount)
|
| },
|
| {ok, State}.
|
|
|
|
|
| -spec handle_call(term(), {pid(), term()}, #state{}) -> {reply, term(), #state{}}.
|
|
|
| handle_call({assign_partition, AgentId, Competency}, _From, State) ->
|
| Partition = compute_partition(AgentId, Competency, State#state.partition_count),
|
|
|
|
|
| Key = {AgentId, Competency},
|
| NewAssignments = maps:put(Key, Partition, State#state.partition_assignments),
|
|
|
|
|
| Load = maps:get(Partition, State#state.partition_loads, 0.0),
|
| NewLoads = maps:put(Partition, Load + 0.01, State#state.partition_loads),
|
|
|
| NewState = State#state{
|
| partition_assignments = NewAssignments,
|
| partition_loads = NewLoads
|
| },
|
|
|
| {reply, Partition, NewState};
|
|
|
| handle_call({get_partition_load, Partition}, _From, State) ->
|
| case maps:find(Partition, State#state.partition_loads) of
|
| {ok, Load} ->
|
| {reply, Load, State};
|
| error ->
|
| {reply, {error, not_found}, State}
|
| end;
|
|
|
| handle_call(rebalance_partitions, _From, State) ->
|
|
|
| MaxLoad = maps:fold(fun(_P, Load, Max) -> max(Load, Max) end, 0.0, State#state.partition_loads),
|
|
|
| case MaxLoad > ?PARTITION_LOAD_THRESHOLD of
|
| true ->
|
|
|
| NewLoads = init_loads(State#state.partition_count),
|
| NewState = State#state{partition_loads = NewLoads},
|
| {reply, ok, NewState};
|
| false ->
|
| {reply, ok, State}
|
| end;
|
|
|
| handle_call(_Request, _From, State) ->
|
| {reply, {error, unknown_call}, State}.
|
|
|
|
|
| -spec handle_cast(term(), #state{}) -> {noreply, #state{}}.
|
| handle_cast(_Msg, State) ->
|
| {noreply, State}.
|
|
|
|
|
| -spec handle_info(term(), #state{}) -> {noreply, #state{}}.
|
| handle_info(_Info, State) ->
|
| {noreply, State}.
|
|
|
|
|
| -spec terminate(term(), #state{}) -> ok.
|
| terminate(_Reason, _State) ->
|
| ok.
|
|
|
|
|
| -spec code_change(term(), #state{}, term()) -> {ok, #state{}}.
|
| code_change(_OldVsn, State, _Extra) ->
|
| {ok, State}.
|
|
|
|
|
|
|
|
|
|
|
|
|
| -spec init_loads(non_neg_integer()) -> map().
|
| init_loads(PartitionCount) ->
|
| maps:from_list([{P, 0.0} || P <- lists:seq(0, PartitionCount - 1)]).
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| -spec compute_partition(binary(), atom(), non_neg_integer()) -> non_neg_integer().
|
| compute_partition(AgentId, Competency, PartitionCount) ->
|
| erlang:phash2({AgentId, Competency}) rem PartitionCount.
|
|
|