Yet Another Riak Client Pool

connection_sup.erl

% connection supervisor
-module (connection_sup).
-behaviour (supervisor).
-export ([start_link/1, start_connection/0, init/1]).

start_link({Ip, Port, Options}) ->
  supervisor:start_link({local, ?MODULE}, ?MODULE, [Ip, Port, Options]).

start_connection () ->
  supervisor:start_child (?MODULE, []).

init([Ip, Port, Options]) ->
  { ok, { { simple_one_for_one, 10, 10 },
          [ { connections,
              { pool_manager, start_connection, [Ip, Port, Options] },
              transient, 2000, worker, dynamic
            }
          ]
    } }.

pool_manager.erl

% pool manager
-module (pool_manager).
-behaviour (gen_server).
-export ([ start_link/1, start_connection/3,
           get_connection/0, all_connections/0,
           init/1, handle_call/3, handle_cast/2, handle_info/2, terminate/2,
           code_change/3 ]).
-define (POOL, connections).
-record (state, {child_max}).

start_link (Config) ->
  gen_server:start_link ( { local, ?MODULE }, ?MODULE, [Config], []).

% connection supervisor calls here so we can startlink and add to the
% pool
start_connection (Ip, Port, Options) ->
  case riakc_pb_socket:start_link (Ip, Port, Options) of
    {ok, P} ->
      case pg2:join (?POOL, P) of
        ok ->
          {ok, P};
        {error, Reason} ->
          riakc_pb_socket:stop(P),
          {error, Reason}
      end;
    E ->
      E
  end.

get_connection () ->
  pg2:get_closest_pid (?POOL).

all_connections () ->
  pg2:get_local_members (?POOL).

% init here just to create the pool and start the connections
init ([ChildMax]) ->
  pg2:create (?POOL),
  % start up the children
  [ {ok, _} = connection_sup:start_connection ()
    || _ <- lists:seq (1, ChildMax) ],
  { ok, #state {child_max = ChildMax} }.

handle_call (_Request, _From, State) -> { reply, ok, State }.
handle_cast (_Request, State) -> { noreply, State }.
handle_info (_Request, State) -> {noreply, State}.
terminate (_Reason, _State) -> ok.
code_change (_OldVsn, State, _Extra) -> {ok, State}.

pool_sup.erl

% top-level supervisor
-module (pool_sup).
-behaviour (supervisor).
-export ([ start_link/1, init/1 ]).
start_link(Config) ->
  supervisor:start_link({local, ?MODULE}, ?MODULE, [Config]).
init([PoolConfig]) ->
  { ok, { { one_for_one, 10, 10 },
          [
            { connection_sup,
              { connection_sup, start_link,
                [ proplists:get_value (connection_config,
                                       PoolConfig,
                                       { "127.0.0.1", 8087, [] }) ]
              }, permanent, 2000, supervisor, [ connection_sup ] },
            { pool_manager,
              { pool_manager, start_link,
                [ proplists:get_value (connection_num,
                                       PoolConfig,
                                       10)
                ]
              }, permanent, 2000, worker, [ pool_manager ]            } ]
   } }.