[#30] – Allow to create the state passing the rest of its attributes

Author: cabol

Makefile

@@ -44,7 +44,7 @@ test: check_rebar check_epmd
 	rm -rf test/*.beam
 
 local_test: check_rebar check_epmd
-	$(REBAR) ct --suite=test/local_SUITE
+	$(REBAR) ct --suite=test/task_SUITE,test/state_SUITE,test/local_SUITE
 	$(REBAR) cover
 	rm -rf test/*.beam
 

README.md

@@ -23,7 +23,7 @@ Therefore, one of the most common and proven strategies to deal with these probl
 
 Here is where **Shards** comes in. **Shards** makes extremely easy achieve all this, with **zero** effort.
 It provides an API compatible with [ETS](http://erlang.org/doc/man/ets.html) – with few exceptions.
-You can find the list of compatible ETS functions that **Shards** provides [HERE](https://github.com/cabol/shards/issues/1).
+You can check the list of compatible ETS functions that **Shards** provides [HERE](https://github.com/cabol/shards/issues/1).
 
 
 ## Usage
@@ -228,9 +228,9 @@ If any microsecond matters to you, you can skip the call to the control ETS tabl
 
  1. The first option is getting the `state`, and passing it as argument. Now the question is:
     how to get the **State**? Well, it's extremely easy, you can get the `state` when you call
-    `shards:new/2` by first time, or you can call `shards:state/1` or `shards_state:get/1`
-    at any time you want, and then it might be stored within the calling process, or wherever
-    you want. E.g.:
+    `shards:new/2` by first time, or you can call `shards:state/1`, `shards_state:get/1` or
+    `shards_state:new/0,1,2,3,4` at any time you want, and then it might be stored within
+    the calling process, or wherever you want. E.g.:
 
     ```erlang
     % take a look at the 2nd element of the returned tuple, that is the state
@@ -247,6 +247,11 @@ If any microsecond matters to you, you can skip the call to the control ETS tabl
     true
     > shards_local:lookup(mytab, 1, State).
     [{1,1}]
+    
+    % in this case, only the n_shards is different from default, so you
+    % can do this:
+    > shards_local:lookup(mytab, 1, shards_state:new(4)).
+    [{1,1}]
     ```
 
  2. The 2nd option is to call `shards_local` directly without the `state`, but this is only

src/shards_state.erl

@@ -10,6 +10,10 @@
 -export([
   get/1,
   new/0,
+  new/1,
+  new/2,
+  new/3,
+  new/4,
   to_map/1,
   from_map/1
 ]).
@@ -112,23 +116,30 @@
 %%% API
 %%%===================================================================
 
-%% @doc
-%% Returns the `state' for the given table `Tab'.
-%% @end
--spec get(Tab :: atom()) -> state().
-get(Tab) when is_atom(Tab) ->
-  case ets:lookup(Tab, state) of
-    [State] -> State;
-    _       -> throw({badarg, Tab})
-  end.
-
-%% @doc
-%% Creates a new `state' with default values.
-%% @end
 -spec new() -> state().
 new() ->
   #state{}.
 
+-spec new(pos_integer()) -> state().
+new(Shards) ->
+  #state{n_shards = Shards}.
+
+-spec new(pos_integer(), module()) -> state().
+new(Shards, Module) ->
+  #state{n_shards = Shards, module = Module}.
+
+-spec new(pos_integer(), module(), pick_fun()) -> state().
+new(Shards, Module, PickShardFun) ->
+  #state{n_shards = Shards, module = Module, pick_shard_fun = PickShardFun}.
+
+-spec new(pos_integer(), module(), pick_fun(), pick_fun()) -> state().
+new(Shards, Module, PickShardFun, PickNodeFun) ->
+  #state{
+    n_shards       = Shards,
+    module         = Module,
+    pick_shard_fun = PickShardFun,
+    pick_node_fun  = PickNodeFun}.
+
 %% @doc
 %% Converts the given `state' into a `map'.
 %% @end
@@ -150,6 +161,16 @@ from_map(Map) ->
     pick_shard_fun = maps:get(pick_shard_fun, Map, fun shards_local:pick/3),
     pick_node_fun  = maps:get(pick_node_fun, Map, fun shards_local:pick/3)}.
 
+%% @doc
+%% Returns the `state' for the given table `Tab'.
+%% @end
+-spec get(Tab :: atom()) -> state().
+get(Tab) when is_atom(Tab) ->
+  case ets:lookup(Tab, state) of
+    [State] -> State;
+    _       -> throw({badarg, Tab})
+  end.
+
 %%%===================================================================
 %%% API – Getters & Setters
 %%%===================================================================
@@ -165,8 +186,8 @@ module(Module, #state{} = State) when is_atom(Module) ->
   State#state{module = Module}.
 
 -spec n_shards(state() | atom()) -> pos_integer().
-n_shards(#state{n_shards = NumShards}) ->
-  NumShards;
+n_shards(#state{n_shards = Shards}) ->
+  Shards;
 n_shards(Tab) when is_atom(Tab) ->
   n_shards(?MODULE:get(Tab)).
 

src/shards_task.erl

@@ -11,15 +11,30 @@
 -module(shards_task).
 
 %% Task API
--export([start/1, start/2, start/3]).
--export([start_link/1, start_link/2, start_link/3]).
--export([async/1, async/2, async/3]).
--export([await/1, await/2]).
+-export([
+  start/1,
+  start/2,
+  start/3,
+  start_link/1,
+  start_link/2,
+  start_link/3,
+  async/1,
+  async/2,
+  async/3,
+  await/1,
+  await/2
+]).
 
 %% Task Supervised API
--export([sup_start/2, sup_start_link/2, sup_start_link/4]).
--export([sup_spawn_link/3, sup_spawn_link/4]).
--export([reply/4, noreply/2]).
+-export([
+  sup_start/2,
+  sup_start_link/2,
+  sup_start_link/4,
+  sup_spawn_link/3,
+  sup_spawn_link/4,
+  reply/4,
+  noreply/2
+]).
 
 -define(TIMEOUT, 5000).
 
@@ -285,8 +300,7 @@ initial_call(MFA) ->
   put('$initial_call', get_initial_call(MFA)).
 
 %% @private
-get_initial_call({erlang, apply, [Fun, Args]})
-    when is_function(Fun, length(Args)) ->
+get_initial_call({erlang, apply, [Fun, Args]}) when is_function(Fun, length(Args)) ->
   {module, Module} = erlang:fun_info(Fun, module),
   {name, Name} = erlang:fun_info(Fun, name),
   {Module, Name, length(Args)};
@@ -313,8 +327,10 @@ exit(Info, MFA, LogReason, Reason) ->
   exit(Reason).
 
 %% @private
-get_from({Node, PidOrName}) when Node == node() -> PidOrName;
-get_from(Other)                                 -> Other.
+get_from({Node, PidOrName}) when Node == node() ->
+  PidOrName;
+get_from(Other) ->
+  Other.
 
 %% @private
 get_running({erlang, apply, [Fun, Args]}) when is_function(Fun, length(Args)) ->
@@ -323,8 +339,7 @@ get_running({Mod, Fun, Args}) ->
   {erlang:make_fun(Mod, Fun, length(Args)), Args}.
 
 %% @private
-get_reason({undef, [{Mod, Fun, Args, _Info} | _] = Stacktrace} = Reason)
-    when is_atom(Mod) and is_atom(Fun) ->
+get_reason({undef, [{Mod, Fun, Args, _Info} | _] = Stacktrace} = Reason) when is_atom(Mod) and is_atom(Fun) ->
   FunExported = fun
     (M, F, A) when is_list(A) ->
       erlang:function_exported(M, F, length(A));
@@ -346,9 +361,13 @@ get_reason(Reason) ->
   Reason.
 
 %% @private
-reason(noconnection, Proc) -> {nodedown, monitor_node(Proc)};
-reason(Reason, _)          -> Reason.
+reason(noconnection, Proc) ->
+  {nodedown, monitor_node(Proc)};
+reason(Reason, _) ->
+  Reason.
 
 %% @private
-monitor_node(Pid) when is_pid(Pid) -> node(Pid);
-monitor_node({_, Node})            -> Node.
+monitor_node(Pid) when is_pid(Pid) ->
+  node(Pid);
+monitor_node({_, Node}) ->
+  Node.

test/dist_SUITE.erl

@@ -16,7 +16,6 @@
 -include_lib("mixer/include/mixer.hrl").
 -mixin([
   {test_helper, [
-    t_state/1,
     t_basic_ops/1
   ]}
 ]).
@@ -44,7 +43,6 @@ all() ->
 groups() ->
   [{dist_test_group, [sequence], [
     t_join_leave_ops,
-    t_state,
     t_basic_ops,
     t_match_ops,
     t_select_ops,

test/local_SUITE.erl

@@ -15,7 +15,6 @@
 -include_lib("mixer/include/mixer.hrl").
 -mixin([
   {test_helper, [
-    t_state/1,
     t_basic_ops/1,
     t_match_ops/1,
     t_select_ops/1,

test/state_SUITE.erl

@@ -0,0 +1,122 @@
+-module(state_SUITE).
+
+-include_lib("common_test/include/ct.hrl").
+
+%% Common Test
+-export([
+  all/0,
+  init_per_suite/1,
+  end_per_suite/1
+]).
+
+%% Test Cases
+-export([
+  t_create_state/1,
+  t_state_ops/1,
+  t_get_state_badarg_error/1
+]).
+
+-define(EXCLUDED_FUNS, [
+  module_info,
+  all,
+  init_per_suite,
+  end_per_suite
+]).
+
+-include("test_helper.hrl").
+
+%%%===================================================================
+%%% Common Test
+%%%===================================================================
+
+all() ->
+  Exports = ?MODULE:module_info(exports),
+  [F || {F, _} <- Exports, not lists:member(F, ?EXCLUDED_FUNS)].
+
+init_per_suite(Config) ->
+  shards:start(),
+  [{scope, l} | Config].
+
+end_per_suite(Config) ->
+  shards:stop(),
+  Config.
+
+%%%===================================================================
+%%% Tests Cases
+%%%===================================================================
+
+t_create_state(_Config) ->
+  % create state with all default attr values
+  State0 = shards_state:new(),
+  shards_local = shards_state:module(State0),
+  true = ?N_SHARDS == shards_state:n_shards(State0),
+  true = fun shards_local:pick/3 == shards_state:pick_shard_fun(State0),
+  true = fun shards_local:pick/3 == shards_state:pick_node_fun(State0),
+
+  % create state using shards_state:new/1
+  State1 = shards_state:new(4),
+  shards_local = shards_state:module(State1),
+  true = 4 == shards_state:n_shards(State1),
+  true = fun shards_local:pick/3 == shards_state:pick_shard_fun(State1),
+  true = fun shards_local:pick/3 == shards_state:pick_node_fun(State1),
+
+  % create state using shards_state:new/2
+  State2 = shards_state:new(2, shards_dist),
+  shards_dist = shards_state:module(State2),
+  true = 2 == shards_state:n_shards(State2),
+  true = fun shards_local:pick/3 == shards_state:pick_shard_fun(State2),
+  true = fun shards_local:pick/3 == shards_state:pick_node_fun(State2),
+
+  % create state using shards_state:new/3
+  Fun = fun(X, Y, Z) -> (X + Y + Z) rem Y end,
+  State3 = shards_state:new(4, shards_dist, Fun),
+  shards_dist = shards_state:module(State3),
+  true = 4 == shards_state:n_shards(State3),
+  true = Fun == shards_state:pick_shard_fun(State3),
+  true = fun shards_local:pick/3 == shards_state:pick_node_fun(State3),
+
+  % create state using shards_state:new/4
+  State4 = shards_state:new(4, shards_dist, Fun, Fun),
+  shards_dist = shards_state:module(State4),
+  true = 4 == shards_state:n_shards(State4),
+  true = Fun == shards_state:pick_shard_fun(State4),
+  true = Fun == shards_state:pick_node_fun(State4),
+
+  ok.
+
+t_state_ops(_Config) ->
+  test_set = shards:new(test_set, [
+    {pick_node_fun, fun test_helper:pick_node/3}
+  ]),
+  StateSet = shards_state:get(test_set),
+  DefaultShards = ?N_SHARDS,
+  #{n_shards := DefaultShards} = shards_state:to_map(StateSet),
+
+  Mod = shards_state:module(test_set),
+  true = Mod == shards_local orelse Mod == shards_dist,
+  DefaultShards = shards_state:n_shards(test_set),
+  Fun1 = fun shards_local:pick/3,
+  Fun1 = shards_state:pick_shard_fun(test_set),
+  Fun2 = fun test_helper:pick_node/3,
+  Fun2 = shards_state:pick_node_fun(test_set),
+
+  State0 = shards_state:new(),
+  State1 = shards_state:module(shards_dist, State0),
+  State2 = shards_state:n_shards(100, State1),
+  Fun = fun(X, Y, Z) -> (X + Y + Z) rem Y end,
+  State3 = shards_state:pick_shard_fun(Fun, State2),
+  State4 = shards_state:pick_node_fun(Fun, State3),
+
+  #{module           := shards_dist,
+    n_shards         := 100,
+    pick_shard_fun   := Fun,
+    pick_node_fun    := Fun
+  } = shards_state:to_map(State4),
+
+  ok.
+
+t_get_state_badarg_error(_Config) ->
+  wrong_tab = ets:new(wrong_tab, [public, named_table]),
+  _ = try shards_state:get(wrong_tab)
+  catch _:{badarg, wrong_tab} -> ok
+  end, ok.