Merge pull request #226 from permaweb/feat/route-providers

feat: dynamic route choice and weighted node selection in `~router@1.0`

Author: samcamwilliams

src/dev_router.erl

@@ -34,7 +34,7 @@
 %% @doc Device function that returns all known routes.
 routes(M1, M2, Opts) ->
     ?event({routes_msg, M1, M2}),
-    Routes = hb_opts:get(routes, [], Opts),
+    Routes = load_routes(Opts),
     ?event({routes, Routes}),
     case hb_ao:get(<<"method">>, M2, Opts) of
         <<"POST">> ->
@@ -75,9 +75,10 @@ routes(M1, M2, Opts) ->
 %% If we have a route that has multiple resolving nodes, check
 %% the load distribution strategy and choose a node. Supported strategies:
 %% <pre>
-%%       All:     Return all nodes (default).
-%%       Random:  Distribute load evenly across all nodes, non-deterministically.
+%%           All: Return all nodes (default).
+%%        Random: Distribute load evenly across all nodes, non-deterministically.
 %%       By-Base: According to the base message's hashpath.
+%%     By-Weight: According to the node's `weight' key.
 %%       Nearest: According to the distance of the node's wallet address to the
 %%                base message's hashpath.
 %% </pre>
@@ -90,7 +91,7 @@ routes(M1, M2, Opts) ->
 %% function, taking only the request message and the `Opts' map.
 route(Msg, Opts) -> route(undefined, Msg, Opts).
 route(_, Msg, Opts) ->
-    Routes = hb_opts:get(routes, [], Opts),
+    Routes = load_routes(Opts),
     R = match_routes(Msg, Routes, Opts),
     ?event({find_route, {msg, Msg}, {routes, Routes}, {res, R}}),
     case (R =/= no_matches) andalso hb_ao:get(<<"node">>, R, Opts) of
@@ -137,6 +138,21 @@ route(_, Msg, Opts) ->
             end
     end.
 
+%% @doc Load the current routes for the node. Allows either explicit routes from
+%% the node message's `routes' key, or dynamic routes generated by resolving the
+%% `routes_provider' message.
+load_routes(Opts) ->
+    case hb_opts:get(routes_provider, not_found, Opts) of
+        not_found -> hb_opts:get(routes, [], Opts);
+        RoutesProvider ->
+            ProviderMsgs = hb_singleton:from(RoutesProvider),
+            ?event({routes_provider, ProviderMsgs}),
+            case hb_ao:resolve_many(ProviderMsgs, Opts) of
+                {ok, Routes} -> Routes;
+                {error, Error} -> throw({routes, routes_provider_failed, Error})
+            end
+    end.
+
 %% @doc Extract the base message ID from a request message. Produces a single
 %% binary ID that can be used for routing decisions.
 extract_base(#{ <<"path">> := Path }, Opts) ->
@@ -231,6 +247,19 @@ choose(0, _, _, _, _) -> [];
 choose(N, <<"Random">>, _, Nodes, _Opts) ->
     Node = lists:nth(rand:uniform(length(Nodes)), Nodes),
     [Node | choose(N - 1, <<"Random">>, nop, lists:delete(Node, Nodes), _Opts)];
+choose(N, <<"By-Weight">>, _, Nodes, Opts) ->
+    NodesWithWeight =
+        [
+            { Node, hb_util:int(hb_ao:get(<<"weight">>, Node, Opts)) }
+        ||
+            Node <- Nodes
+        ],
+    Node = hb_util:weighted_random(NodesWithWeight),
+    [
+        Node
+    |
+        choose(N - 1, <<"By-Weight">>, nop, lists:delete(Node, Nodes), Opts)
+    ];
 choose(N, <<"By-Base">>, Hashpath, Nodes, Opts) when is_binary(Hashpath) ->
     choose(N, <<"By-Base">>, binary_to_bignum(Hashpath), Nodes, Opts);
 choose(N, <<"By-Base">>, HashInt, Nodes, Opts) ->
@@ -303,6 +332,51 @@ binary_to_bignum(Bin) when ?IS_ID(Bin) ->
 
 %%% Tests
 
+routes_provider_test() ->
+    Node = hb_http_server:start_node(#{
+        routes_provider => #{
+            <<"path">> => <<"/test-key/routes">>,
+            <<"test-key">> => #{
+                <<"routes">> => [
+                    #{
+                        <<"template">> => <<"*">>,
+                        <<"node">> => <<"testnode">>
+                    }
+                ]
+            }
+        }
+    }),
+    ?assertEqual(
+        {ok, <<"testnode">>},
+        hb_http:get(Node, <<"/~router@1.0/routes/1/node">>, #{})
+    ).
+
+dynamic_routes_provider_test() ->
+    {ok, Script} = file:read_file("test/test.lua"),
+    Node = hb_http_server:start_node(#{
+        routes_provider => #{
+            <<"device">> => <<"lua@5.3a">>,
+            <<"path">> => <<"routes">>,
+            <<"script">> => Script,
+            <<"node">> => <<"test-dynamic-node">>
+        }
+    }),
+    ?assertEqual(
+        {ok, <<"test-dynamic-node">>},
+        hb_http:get(Node, <<"/~router@1.0/routes/1/node">>, #{})
+    ).
+
+weighted_random_strategy_test() ->
+    Nodes =
+        [
+            #{ <<"host">> => <<"1">>, <<"weight">> => 1 },
+            #{ <<"host">> => <<"2">>, <<"weight">> => 99 }
+        ],
+    SimRes = simulate(1000, 1, Nodes, <<"By-Weight">>),
+    [One, _] = simulation_distribution(SimRes, Nodes),
+    ?assert(One < 20),
+    ?assert(One > 5).
+
 strategy_suite_test_() ->
     lists:map(
         fun(Strategy) ->

src/hb_util.erl

@@ -17,7 +17,7 @@
 -export([print_trace/4, trace_macro_helper/5, print_trace_short/4]).
 -export([ok/1, ok/2, until/1, until/2, until/3]).
 -export([format_trace_short/1, is_hb_module/1, is_hb_module/2, all_hb_modules/0]).
--export([count/2, mean/1, stddev/1, variance/1]).
+-export([count/2, mean/1, stddev/1, variance/1, weighted_random/1]).
 -include("include/hb.hrl").
 
 %%% Simple type coercion functions, useful for quickly turning inputs from the
@@ -709,4 +709,22 @@ stddev(List) ->
 
 variance(List) ->
     Mean = mean(List),
-    lists:sum([ math:pow(X - Mean, 2) || X <- List ]) / length(List).
+    lists:sum([ math:pow(X - Mean, 2) || X <- List ]) / length(List).
+
+%% @doc Shuffle a list.
+shuffle(List) ->
+    [ Y || {_, Y} <- lists:sort([ {rand:uniform(), X} || X <- List]) ].
+
+%% @doc Return a random element from a list, weighted by the values in the list.
+weighted_random(List) ->
+    TotalWeight = lists:sum([ Weight || {_, Weight} <- List ]),
+    Normalized = [ {Item, Weight / TotalWeight} || {Item, Weight} <- List ],
+    Shuffled = shuffle(Normalized),
+    pick_weighted(Shuffled, rand:uniform()).
+
+pick_weighted([], _) ->
+    error(empty_list);
+pick_weighted([{Item, Weight}|_Rest], Remaining) when Remaining < Weight ->
+    Item;
+pick_weighted([{_Item, Weight}|Rest], Remaining) ->
+    pick_weighted(Rest, Remaining - Weight).

test/test.lua

@@ -68,7 +68,24 @@ function preprocess(base, req, opts)
     return { { body = "i like turtles" } }
 end
 
+--- @function sandboxed_fail
+--- @tparam table base
+--- @tparam table request
+--- @error fails when inside the sandbox
 function sandboxed_fail()
     -- Do something that is not dangerous, but is sandboxed nonetheless.
     return os.getenv("PWD")
 end
+
+--- @function routes
+--- @tparam table base
+--- @tparam table request
+--- @return table a table with the `node` field set to the value of the `node`
+--- field in the base message.
+function routes(base, req, opts)
+    return {
+        {
+            node = base.node
+        }
+    }
+end