Add and remove dynamic endpoints

src/grpcbox_channel.erl

@@ -5,6 +5,9 @@
 -export([start_link/3,
          is_ready/1,
          pick/2,
+         pick/3,
+         add_endpoints/2,
+         remove_endpoints/3,
          stop/1,
          stop/2]).
 -export([init/1,
@@ -58,11 +61,19 @@ is_ready(Name) ->
     gen_statem:call(?CHANNEL(Name), is_ready).
 
 %% @doc Picks a subchannel from a pool using the configured strategy.
--spec pick(name(), unary | stream) -> {ok, {pid(), grpcbox_client:interceptor() | undefined}} |
-                                   {error, undefined_channel | no_endpoints}.
+-spec pick(name(), unary | stream) ->
+    {ok, {pid(), grpcbox_client:interceptor() | undefined}} |
+    {error, undefined_channel | no_endpoints}.
 pick(Name, CallType) ->
+    pick(Name, CallType, undefined).
+
+%% @doc Picks a subchannel from a pool using the configured strategy.
+-spec pick(name(), unary | stream, term() | undefined) ->
+    {ok, {pid(), grpcbox_client:interceptor() | undefined}} |
+    {error, undefined_channel | no_endpoints}.
+pick(Name, CallType, Key) ->
     try
-        case gproc_pool:pick_worker(Name) of
+        case pick_worker(Name, Key) of
             false -> {error, no_endpoints};
             Pid when is_pid(Pid) ->
                 {ok, {Pid, interceptor(Name, CallType)}}
@@ -72,6 +83,17 @@ pick(Name, CallType) ->
             {error, undefined_channel}
     end.
 
+pick_worker(Name, undefined) ->
+    gproc_pool:pick_worker(Name);
+pick_worker(Name, Key) ->
+    gproc_pool:pick_worker(Name, Key).
+
+add_endpoints(Name, Endpoints) ->
+    gen_statem:call(?CHANNEL(Name), {add_endpoints, Endpoints}).
+
+remove_endpoints(Name, Endpoints, Reason) ->
+    gen_statem:call(?CHANNEL(Name), {remove_endpoints, Endpoints, Reason}).
+
 -spec interceptor(name(), unary | stream) -> grpcbox_client:interceptor() | undefined.
 interceptor(Name, CallType) ->
     case ets:lookup(?CHANNELS_TAB, {Name, CallType}) of
@@ -100,14 +122,13 @@ init([Name, Endpoints, Options]) ->
         pool = Name,
         encoding = Encoding,
         stats_handler = StatsHandler,
-        endpoints = Endpoints
+        endpoints = lists:usort(Endpoints)
     },
-
     case maps:get(sync_start, Options, false) of
         false ->
             {ok, idle, Data, [{next_event, internal, connect}]};
         true ->
-            _ = start_workers(Name, StatsHandler, Encoding, Endpoints),
+            start_workers(Name, StatsHandler, Encoding, Endpoints),
             {ok, connected, Data}
     end.
 
@@ -116,14 +137,32 @@ callback_mode() ->
 
 connected({call, From}, is_ready, _Data) ->
     {keep_state_and_data, [{reply, From, true}]};
+connected({call, From}, {add_endpoints, Endpoints},
+            Data=#data{pool=Pool,
+                       stats_handler=StatsHandler,
+                       encoding=Encoding,
+                       endpoints=TotalEndpoints}) ->
+    NewEndpoints = lists:subtract(Endpoints, TotalEndpoints),
+    NewTotalEndpoints = lists:umerge(TotalEndpoints, Endpoints),
+    start_workers(Pool, StatsHandler, Encoding, NewEndpoints),
+    {keep_state, Data#data{endpoints=NewTotalEndpoints}, [{reply, From, ok}]};
+connected({call, From}, {remove_endpoints, Endpoints, Reason},
+            Data=#data{pool=Pool, endpoints=TotalEndpoints}) ->
+
+    NewEndpoints = sets:to_list(sets:intersection(sets:from_list(Endpoints),
+                                sets:from_list(TotalEndpoints))),
+    NewTotalEndpoints = lists:subtract(TotalEndpoints, Endpoints),
+    stop_workers(Pool, NewEndpoints, Reason),
+    {keep_state, Data#data{endpoints = NewTotalEndpoints}, [{reply, From, ok}]};
 connected(EventType, EventContent, Data) ->
     handle_event(EventType, EventContent, Data).
 
 idle(internal, connect, Data=#data{pool=Pool,
                                    stats_handler=StatsHandler,
                                    encoding=Encoding,
                                    endpoints=Endpoints}) ->
-    _ = start_workers(Pool, StatsHandler, Encoding, Endpoints),
+
+    start_workers(Pool, StatsHandler, Encoding, Endpoints),
     {next_state, connected, Data};
 idle({call, From}, is_ready, _Data) ->
     {keep_state_and_data, [{reply, From, false}]};
@@ -172,9 +211,16 @@ insert_stream_interceptor(Name, _Type, Interceptors) ->
 
 start_workers(Pool, StatsHandler, Encoding, Endpoints) ->
     [begin
-         gproc_pool:add_worker(Pool, Endpoint),
-         {ok, Pid} = grpcbox_subchannel:start_link(Endpoint, Pool, {Transport, Host, Port, EndpointOptions},
-             Encoding, StatsHandler),
-         Pid
-     end || Endpoint={Transport, Host, Port, EndpointOptions} <- Endpoints].
+        gproc_pool:add_worker(Pool, Endpoint),
+        {ok, Pid} = grpcbox_subchannel:start_link(Endpoint,
+                                                    Pool, Endpoint, Encoding, StatsHandler),
+        Pid
+     end || Endpoint <- Endpoints].
 
+stop_workers(Pool, Endpoints, Reason) ->
+    [begin
+        case gproc_pool:whereis_worker(Pool, Endpoint) of
+            undefined -> ok;
+            Pid -> grpcbox_subchannel:stop(Pid, Reason)
+        end
+     end || Endpoint <- Endpoints].

src/grpcbox_client.erl

@@ -47,7 +47,8 @@
 
 get_channel(Options, Type) ->
     Channel = maps:get(channel, Options, default_channel),
-    grpcbox_channel:pick(Channel, Type).
+    Key =  maps:get(key, Options, undefined),
+    grpcbox_channel:pick(Channel, Type, Key).
 
 unary(Ctx, Service, Method, Input, Def, Options) ->
     unary(Ctx, filename:join([<<>>, Service, Method]), Input, Def, Options).

src/grpcbox_subchannel.erl

@@ -14,7 +14,8 @@
          ready/3,
          disconnected/3]).
 
--record(data, {endpoint :: grpcbox_channel:endpoint(),
+-record(data, {name :: any(),
+               endpoint :: grpcbox_channel:endpoint(),
                channel :: grpcbox_channel:t(),
                info :: #{authority := binary(),
                          scheme := binary(),
@@ -41,8 +42,10 @@ stop(Pid, Reason) ->
 
 init([Name, Channel, Endpoint, Encoding, StatsHandler]) ->
     process_flag(trap_exit, true),
+
     gproc_pool:connect_worker(Channel, Name),
-    {ok, disconnected, #data{conn=undefined,
+    {ok, disconnected, #data{name=Name,
+                             conn=undefined,
                              info=info_map(Endpoint, Encoding, StatsHandler),
                              endpoint=Endpoint,
                              channel=Channel}}.
@@ -89,24 +92,24 @@ handle_event(_, _, _) ->
     keep_state_and_data.
 
 terminate(_Reason, _State, #data{conn=undefined,
-                                 endpoint=Endpoint,
+                                 name=Name,
                                  channel=Channel}) ->
-    gproc_pool:disconnect_worker(Channel, Endpoint),
-    gproc_pool:remove_worker(Channel, Endpoint),
+    gproc_pool:disconnect_worker(Channel, Name),
+    gproc_pool:remove_worker(Channel, Name),
     ok;
 terminate(normal, _State, #data{conn=Pid,
-                                 endpoint=Endpoint,
+                                 name=Name,
                                  channel=Channel}) ->
     h2_connection:stop(Pid),
-    gproc_pool:disconnect_worker(Channel, Endpoint),
-    gproc_pool:remove_worker(Channel, Endpoint),
+    gproc_pool:disconnect_worker(Channel, Name),
+    gproc_pool:remove_worker(Channel, Name),
     ok;
 terminate(Reason, _State, #data{conn=Pid,
-                                 endpoint=Endpoint,
+                                 name=Name,
                                  channel=Channel}) ->
+    gproc_pool:disconnect_worker(Channel, Name),
+    gproc_pool:remove_worker(Channel, Name),
     exit(Pid, Reason),
-    gproc_pool:disconnect_worker(Channel, Endpoint),
-    gproc_pool:remove_worker(Channel, Endpoint),
     ok.
 
 connect(Data=#data{conn=undefined,

test/grpcbox_channel_SUITE.erl

@@ -0,0 +1,39 @@
+-module(grpcbox_channel_SUITE).
+
+-export([all/0,
+        init_per_suite/1,
+        end_per_suite/1,
+        add_and_remove_endpoints/1]).
+
+-include_lib("eunit/include/eunit.hrl").
+
+all() ->
+    [
+        add_and_remove_endpoints
+    ].
+init_per_suite(_Config) ->
+    application:set_env(grpcbox, servers, []),
+    application:ensure_all_started(grpcbox),
+    grpcbox_channel_sup:start_link(),
+    grpcbox_channel_sup:start_child(default_channel, [{http, "127.0.0.1", 18080, []}], #{}),
+    grpcbox_channel_sup:start_child(random_channel,
+                                    [{http, "127.0.0.1", 18080, []}, {http, "127.0.0.1", 18081, []}, {http, "127.0.0.1", 18082, []}, {http, "127.0.0.1", 18083, []}],
+                                    #{balancer => random}),
+    grpcbox_channel_sup:start_child(hash_channel,
+                                    [{http, "127.0.0.1", 18080, []}, {http, "127.0.0.1", 18081, []}, {http, "127.0.0.1", 18082, []}, {http, "127.0.0.1", 18083, []}],
+                                    #{balancer => hash}),
+    grpcbox_channel_sup:start_child(direct_channel,
+                                    [{http, "127.0.0.1", 18080, []}, {http, "127.0.0.1", 18081, []}, {http, "127.0.0.1", 18082, []}, {http, "127.0.0.4", 18084, []}],
+                                    #{ balancer => direct}),
+
+    _Config.
+
+end_per_suite(_Config) ->
+    application:stop(grpcbox),
+    ok.
+
+add_and_remove_endpoints(_Config) ->
+    grpcbox_channel:add_endpoints(default_channel, [{http, "127.0.0.1", 18081, []}, {http, "127.0.0.1", 18082, []}, {http, "127.0.0.1", 18083, []}]),
+    ?assertEqual(4, length(gproc_pool:active_workers(default_channel))),
+    grpcbox_channel:add_endpoints(default_channel, [{https, "127.0.0.1", 18081, []}, {https, "127.0.0.1", 18082, []}, {https, "127.0.0.1", 18083, []}]),
+    ?assertEqual(7, length(gproc_pool:active_workers(default_channel))).