Add multiplexing draft

rebar.config

@@ -3,7 +3,8 @@
             {parse_transform, lager_transform}]}.
 
 {deps, [{ranch, "1.7.0"},
-        {lager, "3.6.3"}
+        {lager, "3.6.3"},
+        {psq, {git, "https://github.com/eryx67/psq.git", {branch, "master"}}}
         ]}.
 
 {xref_checks,

rebar.lock

 {"1.1.0",
 [{<<"goldrush">>,{pkg,<<"goldrush">>,<<"0.1.9">>},1},
  {<<"lager">>,{pkg,<<"lager">>,<<"3.6.3">>},0},
+ {<<"psq">>,
+  {git,"https://github.com/eryx67/psq.git",
+       {ref,"acf8cb6620a9f9cb6123cc45aeb8767fa1a2ab08"}},
+  0},
  {<<"ranch">>,{pkg,<<"ranch">>,<<"1.7.0">>},0}]}.
 [
 {pkg_hash,[

src/mtp_config.erl

@@ -15,21 +15,34 @@
 
 %% API
 -export([start_link/0]).
--export([get_downstream/1,
-         get_downstream_safe/1,
+-export([get_downstream_safe/2,
+         get_downstream_pool/1,
+         get_netloc/1,
+         get_netloc_safe/1,
          get_secret/0]).
+-export([register_name/2,
+         unregister_name/1,
+         whereis_name/1,
+         send/2]).
 
 %% gen_server callbacks
 -export([init/1, handle_call/3, handle_cast/2, handle_info/2,
          terminate/2, code_change/3]).
 
+-type dc_id() :: integer().
+-type netloc() :: {inet:ip4_address(), inet:port_number()}.
+
 -define(TAB, ?MODULE).
+-define(IPS_KEY(DcId), {id, DcId}).
+-define(POOL_KEY(DcId), {pool, DcId}).
+-define(IDS_KEY, dc_ids).
 -define(SECRET_URL, "https://core.telegram.org/getProxySecret").
 -define(CONFIG_URL, "https://core.telegram.org/getProxyConfig").
 
 -define(APP, mtproto_proxy).
 
 -record(state, {tab :: ets:tid(),
+                monitors = #{} :: #{pid() => {reference(), dc_id()}},
                 timer :: gen_timeout:tout()}).
 
 -ifndef(OTP_RELEASE).                           % pre-OTP21
@@ -45,28 +58,74 @@
 start_link() ->
     gen_server:start_link({local, ?MODULE}, ?MODULE, [], []).
 
--spec get_downstream(integer()) -> {ok, {inet:ip4_address(), inet:port_number()}}.
-get_downstream_safe(DcId) ->
-    case get_downstream(DcId) of
-        {ok, Addr} -> Addr;
+-spec get_downstream_safe(dc_id(), mtp_down_conn:upstream_opts()) ->
+                                 {dc_id(), pid(), mtp_down_conn:handle()}.
+get_downstream_safe(DcId, Opts) ->
+    case get_downstream_pool(DcId) of
+        {ok, Pool} ->
+            Downstream = mtp_dc_pool:get(Pool, self(), Opts),
+            {DcId, Pool, Downstream};
         not_found ->
-            [{_, L}] = ets:lookup(?TAB, id_range),
+            [{?IDS_KEY, L}] = ets:lookup(?TAB, ?IDS_KEY),
+            NewDcId = random_choice(L),
+            get_downstream_safe(NewDcId, Opts)
+    end.
+
+get_downstream_pool(DcId) ->
+    Key = ?POOL_KEY(DcId),
+    case ets:lookup(?TAB, Key) of
+        [] -> not_found;
+        [{Key, PoolPid}] ->
+            {ok, PoolPid}
+    end.
+
+-spec get_netloc_safe(dc_id()) -> {dc_id(), netloc()}.
+get_netloc_safe(DcId) ->
+    case get_netloc(DcId) of
+        {ok, Addr} -> {DcId, Addr};
+        not_found ->
+            [{?IDS_KEY, L}] = ets:lookup(?TAB, ?IDS_KEY),
             NewDcId = random_choice(L),
             %% Get random DC; it might return 0 and recurse aggain
-            get_downstream_safe(NewDcId)
+            get_netloc_safe(NewDcId)
     end.
 
-get_downstream(DcId) ->
-    case ets:lookup(?TAB, {id, DcId}) of
+get_netloc(DcId) ->
+    Key = ?IPS_KEY(DcId),
+    case ets:lookup(?TAB, Key) of
         [] ->
             not_found;
-        [{_, Ip, Port}] ->
-            {ok, {Ip, Port}};
-        L ->
-            {_, Ip, Port} = random_choice(L),
-            {ok, {Ip, Port}}
+        [{Key, [{_, _} = IpPort]}] ->
+            {ok, IpPort};
+        [{Key, L}] ->
+            IpPort = random_choice(L),
+            {ok, IpPort}
     end.
 
+register_name(DcId, Pid) ->
+    case ets:insert_new(?TAB, {?POOL_KEY(DcId), Pid}) of
+        true ->
+            gen_server:cast(?MODULE, {reg, DcId, Pid}),
+            yes;
+        false -> no
+    end.
+
+unregister_name(DcId) ->
+    %% making async monitors is a bad idea..
+    Pid = whereis_name(DcId),
+    gen_server:cast(?MODULE, {unreg, DcId, Pid}),
+    ets:delete(?TAB, ?POOL_KEY(DcId)).
+
+whereis_name(DcId) ->
+    case get_downstream_pool(DcId) of
+        not_found -> undefined;
+        {ok, PoolPid} -> PoolPid
+    end.
+
+send(Name, Msg) ->
+    whereis_name(Name) ! Msg.
+
+
 -spec get_secret() -> binary().
 get_secret() ->
     [{_, Key}] = ets:lookup(?TAB, key),
@@ -79,8 +138,8 @@ init([]) ->
     Timer = gen_timeout:new(
               #{timeout => {env, ?APP, conf_refresh_interval, 3600},
                 unit => second}),
-    Tab = ets:new(?TAB, [bag,
-                         protected,
+    Tab = ets:new(?TAB, [set,
+                         public,
                          named_table,
                          {read_concurrency, true}]),
     State = #state{tab = Tab,
@@ -92,8 +151,14 @@ init([]) ->
 handle_call(_Request, _From, State) ->
     Reply = ok,
     {reply, Reply, State}.
-handle_cast(_Msg, State) ->
-    {noreply, State}.
+handle_cast({reg, DcId, Pid}, #state{monitors = Mons} = State) ->
+    Ref = erlang:monitor(process, Pid),
+    Mons1 = Mons#{Pid => {Ref, DcId}},
+    {noreply, State#state{monitors = Mons1}};
+handle_cast({unreg, DcId, Pid}, #state{monitors = Mons} = State) ->
+    {{Ref, DcId}, Mons1} = maps:take(Pid, Mons),
+    erlang:demonitor(Ref, [flush]),
+    {noreply, State#state{monitors = Mons1}}.
 handle_info(timeout, #state{timer = Timer} =State) ->
     case gen_timeout:is_expired(Timer) of
         true ->
@@ -105,8 +170,10 @@ handle_info(timeout, #state{timer = Timer} =State) ->
         false ->
             {noreply, State#state{timer = gen_timeout:reset(Timer)}}
     end;
-handle_info(_Info, State) ->
-    {noreply, State}.
+handle_info({'DOWN', MonRef, process, Pid, _Reason}, #state{monitors = Mons} = State) ->
+    {{MonRef, DcId}, Mons1} = maps:take(Pid, Mons),
+    ets:delete(?TAB, ?POOL_KEY(DcId)),
+    {noreply, State#state{monitors = Mons1}}.
 terminate(_Reason, _State) ->
     ok.
 code_change(_OldVsn, State, _Extra) ->
@@ -117,9 +184,9 @@ code_change(_OldVsn, State, _Extra) ->
 %%%===================================================================
 
 update(#state{tab = Tab}, force) ->
+    update_ip(),
     update_key(Tab),
-    update_config(Tab),
-    update_ip();
+    update_config(Tab);
 update(State, _) ->
     try update(State, force)
     catch ?WITH_STACKTRACE(Class, Reason, Stack)
@@ -135,9 +202,8 @@ update_key(Tab) ->
 update_config(Tab) ->
     {ok, Body} = http_get(?CONFIG_URL),
     Downstreams = parse_config(Body),
-    Range = get_range(Downstreams),
     update_downstreams(Downstreams, Tab),
-    update_range(Range, Tab).
+    update_ids(Downstreams, Tab).
 
 parse_config(Body) ->
     Lines = string:lexemes(Body, "\n"),
@@ -158,15 +224,30 @@ parse_downstream(Line) ->
      IpAddr,
      Port}.
 
-get_range(Downstreams) ->
-    [Id || {Id, _, _} <- Downstreams].
-
 update_downstreams(Downstreams, Tab) ->
-    [true = ets:insert(Tab, {{id, Id}, Ip, Port})
-     || {Id, Ip, Port} <- Downstreams].
+    ByDc = lists:foldl(
+             fun({DcId, Ip, Port}, Acc) ->
+                     Netlocs = maps:get(DcId, Acc, []),
+                     Acc#{DcId => [{Ip, Port} | Netlocs]}
+             end, #{}, Downstreams),
+    [true = ets:insert(Tab, {?IPS_KEY(DcId), Netlocs})
+     || {DcId, Netlocs} <- maps:to_list(ByDc)],
+    lists:foreach(
+      fun(DcId) ->
+              case get_downstream_pool(DcId) of
+                  not_found ->
+                      %% process will be registered asynchronously by
+                      %% gen_server:start_link({via, ..
+                      {ok, _Pid} = mtp_dc_pool_sup:start_pool(DcId);
+                  {ok, _} ->
+                      ok
+              end
+      end,
+      maps:keys(ByDc)).
 
-update_range(Range, Tab) ->
-    true = ets:insert(Tab, {id_range, Range}).
+update_ids(Downstreams, Tab) ->
+    Ids = lists:usort([DcId || {DcId, _, _} <- Downstreams]),
+    true = ets:insert(Tab, {?IDS_KEY, Ids}).
 
 update_ip() ->
     case application:get_env(?APP, ip_lookup_services) of
@@ -201,6 +282,7 @@ random_choice(L) ->
     Idx = rand:uniform(length(L)),
     lists:nth(Idx, L).
 
+
 -ifdef(TEST).
 -include_lib("eunit/include/eunit.hrl").
 

src/mtp_dc_pool.erl

+%%%-------------------------------------------------------------------
+%%% @author Sergey <me@seriyps.ru>
+%%% @copyright (C) 2018, Sergey
+%%% @doc
+%%% Process that manages pool of connections to telegram datacenter
+%%% and is responsible for load-balancing between them
+%%% @end
+%%% TODO: monitoring of DC connections! Make 100% sure they are killed when pool
+%%% is killed. Maybe link?
+%%% Created : 14 Oct 2018 by Sergey <me@seriyps.ru>
+%%%-------------------------------------------------------------------
+-module(mtp_dc_pool).
+
+-behaviour(gen_server).
+
+%% API
+-export([start_link/1,
+         get/3,
+         return/2,
+         add_connection/1,
+         ack_connected/2]).
+
+%% gen_server callbacks
+-export([init/1, handle_call/3, handle_cast/2, handle_info/2,
+         terminate/2, code_change/3]).
+
+-define(SERVER, ?MODULE).
+-define(APP, mtproto_proxy).
+
+-type upstream() :: mtp_handler:handle().
+-type downstream() :: mtp_down_conn:handle().
+-type ds_store() :: psq:psq().
+
+-record(state,
+        {dc_id :: mtp_config:dc_id(),
+         upstreams = #{} :: #{upstream() => downstream()},
+         pending_downstreams = [] :: [pid()],
+         downstreams :: ds_store()
+        }).
+
+%%%===================================================================
+%%% API
+%%%===================================================================
+start_link(DcId) ->
+    gen_server:start_link({via, mtp_config, DcId}, ?MODULE, DcId, []).
+
+get(Pool, Upstream, #{addr := _} = Opts) ->
+    gen_server:call(Pool, {get, Upstream, Opts}).
+
+return(Pool, Upstream) ->
+    gen_server:cast(Pool, {return, Upstream}).
+
+add_connection(Pool) ->
+    gen_server:call(Pool, add_connection, 10000).
+
+ack_connected(Pool, Downstream) ->
+    gen_server:cast(Pool, {connected, Downstream}).
+
+%%%===================================================================
+%%% gen_server callbacks
+%%%===================================================================
+init(DcId) ->
+    InitConnections = application:get_env(mtproto_proxy, init_dc_connections, 4),
+    PendingConnections = [do_connect(DcId) || _ <- lists:seq(1, InitConnections)],
+    Connections = recv_pending(PendingConnections),
+    Downstreams = ds_new(Connections),
+    {ok, #state{dc_id = DcId, downstreams = Downstreams}}.
+
+handle_call({get, Upstream, Opts}, _From, State) ->
+    {Downstream, State1} = handle_get(Upstream, Opts, State),
+    {reply, Downstream, State1};
+handle_call(add_connection, _From, State) ->
+    State1 = connect(State),
+    {reply, ok, State1}.
+
+handle_cast({return, Upstream}, State) ->
+    {noreply, handle_return(Upstream, State)};
+handle_cast({connected, Pid}, State) ->
+    {noreply, handle_connected(Pid, State)}.
+
+handle_info({'DOWN', MonitorRef, process, Pid, _Reason}, State) ->
+    %% TODO: monitor downstream connections as well
+    {noreply, handle_down(MonitorRef, Pid, State)}.
+terminate(_Reason, #state{downstreams = Ds}) ->
+    ds_foreach(
+      fun(Pid) ->
+              mtp_down_conn:shutdown(Pid)
+      end, Ds),
+    %% upstreams will be killed by connection itself
+    ok.
+code_change(_OldVsn, State, _Extra) ->
+    {ok, State}.
+
+%%%===================================================================
+%%% Internal functions
+%%%===================================================================
+
+%% Handle async connection ack
+handle_connected(Pid, #state{pending_downstreams = Pending,
+                             downstreams = Ds} = St) ->
+    Pending1 = lists:delete(Pid, Pending),
+    Downstreams1 = ds_add_downstream(Pid, Ds),
+    St#state{pending_downstreams = Pending1,
+             downstreams = Downstreams1}.
+
+handle_get(Upstream, Opts, #state{downstreams = Ds,
+                                  upstreams = Us} = St) ->
+    {Downstream, N, Ds1} = ds_get(Ds),
+    MonRef = erlang:monitor(process, Upstream),
+    %% if N > X and len(pending) < Y -> connect()
+    Us1 = Us#{Upstream => {Downstream, MonRef}},
+    ok = mtp_down_conn:upstream_new(Downstream, Upstream, Opts),
+    {Downstream, maybe_spawn_connection(
+                   N,
+                   St#state{downstreams = Ds1,
+                            upstreams = Us1})}.
+
+handle_return(Upstream, #state{downstreams = Ds,
+                               upstreams = Us} = St) ->
+    {{Downstream, MonRef}, Us1} = maps:take(Upstream, Us),
+    ok = mtp_down_conn:upstream_closed(Downstream, Upstream),
+    erlang:demonitor(MonRef, [flush]),
+    Ds1 = ds_return(Downstream, Ds),
+    St#state{downstreams = Ds1,
+             upstreams = Us1}.
+
+handle_down(MonRef, MaybeUpstream, #state{downstreams = Ds,
+                                          upstreams = Us} = St) ->
+    case maps:take(MaybeUpstream, Us) of
+        {{Downstream, MonRef}, Us1} ->
+            ok = mtp_down_conn:upstream_closed(Downstream, MaybeUpstream),
+            Ds1 = ds_return(Downstream, Ds),
+            St#state{downstreams = Ds1,
+                     upstreams = Us1};
+        error ->
+            lager:warning("Unexpected DOWN. ref=~p, pid=~p", [MonRef, MaybeUpstream]),
+            St
+    end.
+
+maybe_spawn_connection(CurrentMin, #state{pending_downstreams = Pending} = St) ->
+    %% TODO: shrinking (by timer)
+    case application:get_env(?APP, clients_per_dc_connection) of
+        {ok, N} when CurrentMin > N,
+                     Pending == [] ->
+            ToSpawn = 2,
+            lists:foldl(
+              fun(_, S) ->
+                      connect(S)
+              end, St, lists:seq(1, ToSpawn));
+        _ ->
+            St
+    end.
+
+%% Initiate new async connection
+connect(#state{pending_downstreams = Pending,
+               dc_id = DcId} = St) ->
+    %% Should monitor connection PIDs as well!
+    Pid = do_connect(DcId),
+    St#state{pending_downstreams = [Pid | Pending]}.
+
+%% Asynchronous connect
+do_connect(DcId) ->
+    {ok, Pid} = mtp_down_conn_sup:start_conn(self(), DcId),
+    Pid.
+
+%% Block until all async connections are acked
+recv_pending(Pids) ->
+    [receive
+         {'$gen_cast', {connected, Pid}} -> Pid
+     after 10000 ->
+             exit({timeout, receive Smth -> Smth after 0 -> none end})
+     end || Pid <- Pids].
+
+%% New downstream connection storage
+-spec ds_new([downstream()]) -> ds_store().
+ds_new(Connections) ->
+    Psq = pid_psq:new(),
+    %% TODO: add `from_list` function
+    lists:foldl(
+      fun(Conn, Psq1) ->
+              pid_psq:add(Conn, Psq1)
+      end, Psq, Connections).
+
+-spec ds_foreach(fun( (downstream()) -> any() ), ds_store()) -> ok.
+ds_foreach(Fun, St) ->
+    psq:fold(
+      fun(_, _N, Pid, _) ->
+              Fun(Pid)
+      end, ok, St).
+
+%% Add new downstream to storage
+-spec ds_add_downstream(downstream(), ds_store()) -> ds_store().
+ds_add_downstream(Conn, St) ->
+    pid_psq:add(Conn, St).
+
+%% Get least loaded downstream connection
+-spec ds_get(ds_store()) -> {downstream(), pos_integer(), ds_store()}.
+ds_get(St) ->
+    %% TODO: should return real number of connections
+    {ok, {{Conn, N}, St1}} = pid_psq:get_min_priority(St),
+    {Conn, N, St1}.
+
+%% Return connection back to storage
+-spec ds_return(downstream(), ds_store()) -> ds_store().
+ds_return(Pid, St) ->
+    {ok, St1} = pid_psq:dec_priority(Pid, St),
+    St1.

src/mtp_dc_pool_sup.erl

+%%%-------------------------------------------------------------------
+%%% @author Sergey <me@seriyps.ru>
+%%% @copyright (C) 2018, Sergey
+%%% @doc
+%%% Supervisor for mtp_dc_pool processes
+%%% @end
+%%% Created : 14 Oct 2018 by Sergey <me@seriyps.ru>
+%%%-------------------------------------------------------------------
+-module(mtp_dc_pool_sup).
+
+-behaviour(supervisor).
+
+-export([start_link/0,
+         start_pool/1]).
+-export([init/1]).
+
+-define(SERVER, ?MODULE).
+
+start_link() ->
+    supervisor:start_link({local, ?SERVER}, ?MODULE, []).
+
+-spec start_pool(mtp_config:dc_id()) -> {ok, pid()}.
+start_pool(DcId) ->
+    %% Or maybe it should read IPs from mtp_config by itself?
+    supervisor:start_child(?SERVER, [DcId]).
+
+init([]) ->
+
+    SupFlags = #{strategy => simple_one_for_one,
+                 intensity => 50,
+                 period => 5},
+
+    AChild = #{id => mtp_dc_pool,
+               start => {mtp_dc_pool, start_link, []},
+               restart => permanent,
+               shutdown => 10000,
+               type => worker},
+
+    {ok, {SupFlags, [AChild]}}.

src/mtp_down_conn_sup.erl

+%%%-------------------------------------------------------------------
+%%% @author Sergey <me@seriyps.ru>
+%%% @copyright (C) 2018, Sergey
+%%% @doc
+%%% Supervisor for mtp_down_conn processes
+%%% @end
+%%% Created : 14 Oct 2018 by Sergey <me@seriyps.ru>
+%%%-------------------------------------------------------------------
+-module(mtp_down_conn_sup).
+
+-behaviour(supervisor).
+
+-export([start_link/0,
+         start_conn/2]).
+-export([init/1]).
+
+-define(SERVER, ?MODULE).
+
+start_link() ->
+    supervisor:start_link({local, ?SERVER}, ?MODULE, []).
+
+-spec start_conn(pid(), mtp_conf:dc_id()) -> {ok, pid()}.
+start_conn(Pool, DcId) ->
+    supervisor:start_child(?SERVER, [Pool, DcId]).
+
+init([]) ->
+
+    SupFlags = #{strategy => simple_one_for_one,
+                 intensity => 50,
+                 period => 5},
+
+    AChild = #{id => mtp_down_conn,
+               start => {mtp_down_conn, start_link, []},
+               restart => temporary,
+               shutdown => 2000,
+               type => worker},
+
+    {ok, {SupFlags, [AChild]}}.

src/mtproto_proxy.app.src

@@ -6,6 +6,7 @@
   {applications,
    [lager,
     ranch,
+    psq,
     crypto,
     ssl,
     inets,
@@ -55,9 +56,12 @@
         %% only `{allowed_protocols, [mtp_secure]}` if you want to only allow
         %% connections to this proxy with "dd"-secrets. Connections by other
         %% protocols will be immediately closed.
-        {allowed_protocols, [mtp_abridged, mtp_intermediate, mtp_secure]}
+        {allowed_protocols, [mtp_abridged, mtp_intermediate, mtp_secure]},
 
-        %% module with function `notify/4' exported.
+        {init_dc_connections, 2},
+        {clients_per_dc_connection, 300}
+
+        %% Should be module with function `notify/4' exported.
         %% See mtp_metric:notify/4 for details
         %% {metric_backend, my_metric_backend}
 

src/mtproto_proxy_sup.erl

 %%%-------------------------------------------------------------------
 %% @doc mtproto_proxy top level supervisor.
 %% @end
+%% <pre>
+%% dc_pool_sup (simple_one_for_one)
+%%   dc_pool_1 [conn1, conn3, conn4, ..]
+%%   dc_pool_-1 [conn2, ..]
+%%   dc_pool_2 [conn5, conn7, ..]
+%%   dc_pool_-2 [conn6, conn8, ..]
+%%   ...
+%% down_conn_sup (simple_one_for_one)
+%%   conn1
+%%   conn2
+%%   conn3
+%%   conn4
+%%   ...
+%%   connN
+%% </pre>
 %%%-------------------------------------------------------------------
 
 -module(mtproto_proxy_sup).
@@ -26,16 +41,17 @@ start_link() ->
 %% Supervisor callbacks
 %%====================================================================
 
-%% Child :: {Id,StartFunc,Restart,Shutdown,Type,Modules}
 init([]) ->
-    Childs = [#{id => mtp_config,
+    SupFlags = #{strategy => one_for_all,       %TODO: maybe change strategy
+                 intensity => 50,
+                 period => 5},
+    Childs = [#{id => mtp_down_conn_sup,
+                type => supervisor,
+                start => {mtp_down_conn_sup, start_link, []}},
+              #{id => mtp_dc_pool_sup,
+                type => supervisor,
+                start => {mtp_dc_pool_sup, start_link, []}},
+              #{id => mtp_config,
                 start => {mtp_config, start_link, []}}
              ],
-    {ok, {#{strategy => rest_for_one,
-            intensity => 50,
-            period => 5},
-          Childs} }.
-
-%%====================================================================
-%% Internal functions
-%%====================================================================
+    {ok, {SupFlags, Childs}}.