├── .gitignore ├── Makefile ├── README.markdown ├── rebar.config └── src ├── riakpool.app.src ├── riakpool.erl ├── riakpool_app.erl ├── riakpool_client.erl ├── riakpool_connection_sup.erl └── riakpool_sup.erl /.gitignore: -------------------------------------------------------------------------------- 1 | ebin 2 | doc 3 | deps 4 | .eunit 5 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | REBAR := rebar 2 | 3 | .PHONY: all deps doc test clean update release 4 | 5 | all: deps 6 | $(REBAR) compile 7 | 8 | deps: 9 | $(REBAR) get-deps 10 | 11 | doc: 12 | $(REBAR) doc skip_deps=true 13 | 14 | test: 15 | $(REBAR) eunit skip_deps=true 16 | 17 | clean: 18 | $(REBAR) clean 19 | 20 | update: 21 | git pull 22 | $(REBAR) update-deps 23 | $(REBAR) compile 24 | 25 | release: all test 26 | dialyzer --src src/*.erl 27 | -------------------------------------------------------------------------------- /README.markdown: -------------------------------------------------------------------------------- 1 |

2 | 3 |

4 | 5 | Overview 6 | -------- 7 | riakpool is an application for maintaining a dynamic pool of protocol buffer 8 | client connections to a riak database. It ensures that a given connection can 9 | only be in use by one external process at a time. 10 | 11 | Installation 12 | ------------ 13 | $ git clone git://github.com/dweldon/riakpool.git 14 | $ cd riakpool 15 | $ make 16 | 17 | Interface 18 | --------- 19 | The following example gives an overview of the riakpool interface. Please see 20 | the complete documentation by running `make doc`. 21 | 22 | 1> application:start(riakpool). 23 | ok 24 | 2> riakpool:start_pool("127.0.0.1", 8087). 25 | ok 26 | 3> riakpool:execute(fun(C) -> riakc_pb_socket:ping(C) end). 27 | {ok,pong} 28 | 4> riakpool_client:put(<<"groceries">>, <<"mine">>, <<"eggs">>). 29 | ok 30 | 5> riakpool_client:get(<<"groceries">>, <<"mine">>). 31 | {ok,<<"eggs">>} 32 | 6> riakpool_client:list_keys(<<"groceries">>). 33 | {ok,[<<"mine">>]} 34 | 7> riakpool_client:delete(<<"groceries">>, <<"mine">>). 35 | ok 36 | 8> riakpool:count(). 37 | 1 38 | 39 | Note that the use of riakpool_client is completely optional - it is simply a 40 | collection of convenience functions which call riakpool:execute/1. 41 | 42 | Starting the Pool 43 | ----------------- 44 | Prior to any calls to `riakpool:execute/1`, the pool must be started. This can 45 | be accomplished in one of two ways: 46 | 47 | 1. Before the server is started, set the riakpool application environment 48 | variables `riakpool_host` and `riakpool_port`. 49 | 2. After the server is started, call `riakpool:start_pool/0` or 50 | `riakpool:start_pool/2` (see previous section). 51 | -------------------------------------------------------------------------------- /rebar.config: -------------------------------------------------------------------------------- 1 | %%-*- mode: erlang -*- 2 | {cover_enabled, true}. 3 | {deps, [{riakc, ".", {git, "git://github.com/basho/riak-erlang-client.git", "HEAD"}}]}. 4 | -------------------------------------------------------------------------------- /src/riakpool.app.src: -------------------------------------------------------------------------------- 1 | {application, riakpool, 2 | [{description, "riakpool"}, 3 | {vsn, "0.1"}, 4 | {registered, []}, 5 | {applications, [kernel, stdlib]}, 6 | {mod, {riakpool_app, []}}, 7 | {env, []}]}. 8 | -------------------------------------------------------------------------------- /src/riakpool.erl: -------------------------------------------------------------------------------- 1 | %% @author David Weldon 2 | %% @doc riakpool implements a pool of riak protocol buffer clients. In order to 3 | %% use a connection, a call to {@link execute/1} must be made. This will check 4 | %% out a connection from the pool, use it, and check it back in. This ensures 5 | %% that a given connection can only be in use by one external process at a time. 6 | %% If no existing connections are found, a new one will be established. Note 7 | %% this means the pool will always be the size of the last peak need. The number 8 | %% of connections can be checked with {@link count/0}. 9 | %% 10 | %% Prior to any calls to {@link execute/1}, the pool must be started. This can 11 | %% be accomplished in one of two ways: 12 | %% 13 | %% 1. Before the server is started, set the riakpool application environment 14 | %% variables `riakpool_host' and `riakpool_port'. 15 | %% 16 | %% 2. After the server is started, call {@link start_pool/0} or 17 | %% {@link start_pool/2} 18 | 19 | -module(riakpool). 20 | -behaviour(gen_server). 21 | -export([count/0, 22 | execute/1, 23 | start_link/0, 24 | start_pool/0, 25 | start_pool/2, 26 | stop/0]). 27 | -export([init/1, 28 | handle_call/3, 29 | handle_cast/2, 30 | handle_info/2, 31 | terminate/2, 32 | code_change/3]). 33 | 34 | -type host() :: string() | atom(). 35 | 36 | -record(state, {host :: host(), port ::non_neg_integer(), pids}). 37 | 38 | %% @doc Returns the number of connections as seen by the supervisor. 39 | -spec count() -> integer(). 40 | count() -> 41 | Props = supervisor:count_children(riakpool_connection_sup), 42 | case proplists:get_value(active, Props) of 43 | N when is_integer(N) -> N; 44 | undefined -> 0 45 | end. 46 | 47 | %% @doc Finds the next available connection pid from the pool and calls 48 | %% `Fun(Pid)'. Returns `{ok, Value}' if the call was successful, and 49 | %% `{error, any()}' otherwise. If no connection could be found, a new connection 50 | %% will be established. 51 | %% ``` 52 | %% > riakpool:execute(fun(C) -> riakc_pb_socket:ping(C) end). 53 | %% {ok,pong} 54 | %% ''' 55 | -spec execute(fun((pid()) -> any())) -> {ok, Value::any()} | {error, any()}. 56 | execute(Fun) -> 57 | case gen_server:call(?MODULE, check_out) of 58 | {ok, Pid} -> 59 | try {ok, Fun(Pid)} 60 | catch _:E -> {error, E} 61 | after gen_server:cast(?MODULE, {check_in, Pid}) end; 62 | {error, E} -> {error, E} 63 | end. 64 | 65 | %% @doc Starts the server. 66 | -spec start_link() -> {ok, pid()} | {error, any()}. 67 | start_link() -> gen_server:start_link({local, ?MODULE}, ?MODULE, [], []). 68 | 69 | %% @doc Starts a connection pool to a server listening on {"127.0.0.1", 8087}. 70 | %% Note that a pool can only be started once. 71 | -spec start_pool() -> ok | {error, any()}. 72 | start_pool() -> start_pool("127.0.0.1", 8087). 73 | 74 | %% @doc Starts a connection pool to a server listening on {`Host', `Port'}. 75 | %% Note that a pool can only be started once. 76 | -spec start_pool(host(), integer()) -> ok | {error, any()}. 77 | start_pool(Host, Port) when is_integer(Port) -> 78 | gen_server:call(?MODULE, {start_pool, Host, Port}). 79 | 80 | %% @doc Stops the server. 81 | -spec stop() -> ok. 82 | stop() -> gen_server:cast(?MODULE, stop). 83 | 84 | %% @hidden 85 | init([]) -> 86 | process_flag(trap_exit, true), 87 | case [application:get_env(P) || P <- [riakpool_host, riakpool_port]] of 88 | [{ok, Host}, {ok, Port}] when is_integer(Port) -> 89 | {ok, new_state(Host, Port)}; 90 | _ -> {ok, undefined} 91 | end. 92 | 93 | %% @hidden 94 | handle_call({start_pool, Host, Port}, _From, undefined) -> 95 | case new_state(Host, Port) of 96 | State=#state{} -> {reply, ok, State}; 97 | undefined -> {reply, {error, connection_error}, undefined} 98 | end; 99 | handle_call({start_pool, _Host, _Port}, _From, State=#state{}) -> 100 | {reply, {error, pool_already_started}, State}; 101 | handle_call(check_out, _From, undefined) -> 102 | {reply, {error, pool_not_started}, undefined}; 103 | handle_call(check_out, _From, State=#state{host=Host, port=Port, pids=Pids}) -> 104 | case next_pid(Host, Port, Pids) of 105 | {ok, Pid, NewPids} -> {reply, {ok, Pid}, State#state{pids=NewPids}}; 106 | {error, NewPids} -> 107 | {reply, {error, connection_error}, State#state{pids=NewPids}} 108 | end; 109 | handle_call(_Request, _From, State) -> {reply, ok, State}. 110 | 111 | %% @hidden 112 | handle_cast({check_in, Pid}, State=#state{pids=Pids}) -> 113 | NewPids = queue:in(Pid, Pids), 114 | {noreply, State#state{pids=NewPids}}; 115 | handle_cast(stop, State) -> {stop, normal, State}; 116 | handle_cast(_Msg, State) -> {noreply, State}. 117 | 118 | %% @hidden 119 | handle_info(_Info, State) -> {noreply, State}. 120 | 121 | %% @hidden 122 | terminate(_Reason, undefined) -> ok; 123 | terminate(_Reason, #state{pids=Pids}) -> 124 | StopFun = 125 | fun(Pid) -> 126 | case is_process_alive(Pid) of 127 | true -> riakc_pb_socket:stop(Pid); 128 | false -> ok 129 | end 130 | end, 131 | [StopFun(Pid) || Pid <- queue:to_list(Pids)], ok. 132 | 133 | %% @hidden 134 | code_change(_OldVsn, State, _Extra) -> {ok, State}. 135 | 136 | %% @doc Returns a state with a single pid if a connection could be established, 137 | %% otherwise returns undefined. 138 | -spec new_state(host(), integer()) -> #state{} | undefined. 139 | new_state(Host, Port) -> 140 | case new_connection(Host, Port) of 141 | {ok, Pid} -> 142 | #state{host=Host, port=Port, pids=queue:in(Pid, queue:new())}; 143 | error -> undefined 144 | end. 145 | 146 | %% @doc Returns {ok, Pid} if a new connection was established and added to the 147 | %% supervisor, otherwise returns error. 148 | -spec new_connection(host(), integer()) -> {ok, pid()} | error. 149 | new_connection(Host, Port) -> 150 | case supervisor:start_child(riakpool_connection_sup, [Host, Port]) of 151 | {ok, Pid} when is_pid(Pid) -> {ok, Pid}; 152 | {ok, Pid, _} when is_pid(Pid) -> {ok, Pid}; 153 | _ -> error 154 | end. 155 | 156 | %% @doc Recursively dequeues Pids in search of a live connection. Dead 157 | %% connections are removed from the queue as it is searched. If no connection 158 | %% pid could be found, a new one will be established. Returns {ok, Pid, NewPids} 159 | %% where NewPids is the queue after any necessary dequeues. Returns error if no 160 | %% live connection could be found and no new connection could be established. 161 | -spec next_pid(host(), integer(), queue:queue()) -> {ok, pid(), queue:queue()} | 162 | {error, queue:queue()}. 163 | next_pid(Host, Port, Pids) -> 164 | case queue:out(Pids) of 165 | {{value, Pid}, NewPids} -> 166 | case is_process_alive(Pid) of 167 | true -> {ok, Pid, NewPids}; 168 | false -> next_pid(Host, Port, NewPids) 169 | end; 170 | {empty, _} -> 171 | case new_connection(Host, Port) of 172 | {ok, Pid} -> {ok, Pid, Pids}; 173 | error -> {error, Pids} 174 | end 175 | end. 176 | 177 | -ifdef(TEST). 178 | -include_lib("eunit/include/eunit.hrl"). 179 | 180 | execute_test() -> 181 | riakpool_connection_sup:start_link(), 182 | riakpool:start_link(), 183 | riakpool:start_pool(), 184 | ?assertEqual(1, count()), 185 | Fun1 = fun(C) -> riakc_pb_socket:ping(C) end, 186 | Fun2 = fun(_) -> riakc_pb_socket:ping(1) end, 187 | ?assertEqual({ok, pong}, execute(Fun1)), 188 | ?assertMatch({error, _}, execute(Fun2)), 189 | ?assertEqual({ok, pong}, execute(Fun1)), 190 | riakpool:stop(), 191 | timer:sleep(10), 192 | ?assertEqual(0, count()). 193 | 194 | execute_error_test() -> 195 | riakpool:start_link(), 196 | Fun = fun(C) -> riakc_pb_socket:ping(C) end, 197 | ?assertEqual({error, pool_not_started}, execute(Fun)), 198 | riakpool:stop(), 199 | timer:sleep(10), 200 | ?assertEqual(0, count()). 201 | 202 | start_pool_test() -> 203 | riakpool_connection_sup:start_link(), 204 | riakpool:start_link(), 205 | {H, P} = {"localhost", 8000}, 206 | ?assertEqual({error, connection_error}, riakpool:start_pool(H, P)), 207 | ?assertEqual(ok, riakpool:start_pool()), 208 | ?assertEqual({error, pool_already_started}, riakpool:start_pool()), 209 | riakpool:stop(), 210 | timer:sleep(10), 211 | ?assertEqual(0, count()). 212 | 213 | next_pid_test() -> 214 | riakpool_connection_sup:start_link(), 215 | {H, P} = {"localhost", 8087}, 216 | ?assertEqual(0, count()), 217 | {ok, P1} = new_connection(H, P), 218 | {ok, P2} = new_connection(H, P), 219 | {ok, P3} = new_connection(H, P), 220 | ?assertEqual(3, count()), 221 | riakc_pb_socket:stop(P1), 222 | riakc_pb_socket:stop(P2), 223 | ?assertEqual(1, count()), 224 | Q0 = queue:new(), 225 | Q = queue:from_list([P1, P2, P3]), 226 | ?assertMatch({ok, P3, Q0}, next_pid(H, P, Q)), 227 | riakc_pb_socket:stop(P3), 228 | {ok, P4, Q0} = next_pid(H, P, Q0), 229 | ?assertEqual(1, count()), 230 | riakc_pb_socket:stop(P4), 231 | ?assertEqual(0, count()). 232 | 233 | next_pid_error_test() -> 234 | {H, P} = {"localhost", 8000}, 235 | Q0 = queue:new(), 236 | ?assertMatch({error, Q0}, next_pid(H, P, Q0)). 237 | 238 | -endif. 239 | -------------------------------------------------------------------------------- /src/riakpool_app.erl: -------------------------------------------------------------------------------- 1 | %% @author David Weldon 2 | %% @hidden 3 | 4 | -module(riakpool_app). 5 | -behaviour(application). 6 | -export([start/2, stop/1]). 7 | 8 | start(_StartType, _StartArgs) -> riakpool_sup:start_link(). 9 | 10 | stop(_State) -> ok. 11 | 12 | -ifdef(TEST). 13 | -include_lib("eunit/include/eunit.hrl"). 14 | 15 | app_test() -> 16 | application:set_env(riakpool, riakpool_host, "localhost"), 17 | application:set_env(riakpool, riakpool_port, 8087), 18 | application:start(riakpool), 19 | Fun = fun(C) -> riakc_pb_socket:ping(C) end, 20 | ?assertEqual({ok, pong}, riakpool:execute(Fun)), 21 | ?assertEqual(1, riakpool:count()), 22 | application:stop(riakpool), 23 | application:unset_env(riakpool, riakpool_host), 24 | application:unset_env(riakpool, riakpool_port). 25 | 26 | -endif. 27 | -------------------------------------------------------------------------------- /src/riakpool_client.erl: -------------------------------------------------------------------------------- 1 | %% @author David Weldon 2 | %% @doc riakpool_client is a collection of convenience functions for using 3 | %% riakpool. 4 | 5 | -module(riakpool_client). 6 | -export([delete/2, get/2, list_keys/1, post/3, put/3]). 7 | 8 | %% @doc Delete `Key' from `Bucket'. 9 | -spec delete(binary(), binary()) -> ok. 10 | delete(Bucket, Key) -> 11 | riakpool:execute(fun(C) -> riakc_pb_socket:delete(C, Bucket, Key) end), ok. 12 | 13 | %% @doc Returns the value associated with `Key' in `Bucket' as `{ok, binary()}'. 14 | %% If an error was encountered or the value was not present, returns 15 | %% `{error, any()}'. 16 | -spec get(binary(), binary()) -> {ok, binary()} | {error, any()}. 17 | get(Bucket, Key) -> 18 | Fun = 19 | fun(C) -> 20 | case riakc_pb_socket:get(C, Bucket, Key) of 21 | {ok, O} -> riakc_obj:get_value(O); 22 | {error, E} -> {error, E} 23 | end 24 | end, 25 | case riakpool:execute(Fun) of 26 | {ok, Value} when is_binary(Value) -> {ok, Value}; 27 | {ok, {error, E}} -> {error, E}; 28 | {error, E} -> {error, E} 29 | end. 30 | 31 | %% @doc Returns the list of keys in `Bucket' as `{ok, list()}'. If an error was 32 | %% encountered, returns `{error, any()}'. 33 | -spec list_keys(binary()) -> {ok, list()} | {error, any()}. 34 | list_keys(Bucket) -> 35 | Fun = fun(C) -> riakc_pb_socket:list_keys(C, Bucket) end, 36 | case riakpool:execute(Fun) of 37 | {ok, {ok, Keys}} -> {ok, Keys}; 38 | {error, E} -> {error, E} 39 | end. 40 | 41 | %% @doc Associates `Key' with `Value' in `Bucket' without checking if `Key' 42 | %% already exists in `Bucket'. 43 | -spec post(binary(), binary(), binary()) -> ok. 44 | post(Bucket, Key, Value) -> 45 | Fun = 46 | fun(C) -> 47 | Object = riakc_obj:new(Bucket, Key, Value), 48 | riakc_pb_socket:put(C, Object) 49 | end, 50 | riakpool:execute(Fun), ok. 51 | 52 | %% @doc Associates `Key' with `Value' in `Bucket'. If `Key' already exists in 53 | %% `Bucket', an update will be preformed. 54 | -spec put(binary(), binary(), binary()) -> ok. 55 | put(Bucket, Key, Value) -> 56 | Fun = 57 | fun(C) -> 58 | case riakc_pb_socket:get(C, Bucket, Key) of 59 | {ok, O} -> 60 | O2 = riakc_obj:update_value(O, Value), 61 | riakc_pb_socket:put(C, O2); 62 | {error, _} -> 63 | O = riakc_obj:new(Bucket, Key, Value), 64 | riakc_pb_socket:put(C, O) 65 | end 66 | end, 67 | riakpool:execute(Fun), ok. 68 | 69 | -ifdef(TEST). 70 | -include_lib("eunit/include/eunit.hrl"). 71 | 72 | client_test() -> 73 | {B, K, V1, V2} = {<<"groceries">>, <<"mine">>, <<"eggs">>, <<"toast">>}, 74 | 75 | %% start the application 76 | application:start(riakpool), 77 | 78 | %% nothing should work until the pool is started 79 | ?assertMatch({error, _}, list_keys(B)), 80 | ?assertMatch({error, _}, get(B, K)), 81 | 82 | %% start the pool 83 | riakpool:start_pool(), 84 | 85 | %% verify nothing is in B 86 | ?assertEqual({ok, []}, list_keys(B)), 87 | ?assertMatch({error, notfound}, get(B, K)), 88 | 89 | %% post {K, V1} 90 | ?assertEqual(ok, post(B, K, V1)), 91 | ?assertEqual({ok, V1}, get(B, K)), 92 | 93 | %% delete K 94 | ?assertEqual(ok, delete(B, K)), 95 | ?assertMatch({error, notfound}, get(B, K)), 96 | 97 | %% put {K, V1} 98 | ?assertEqual(ok, put(B, K, V1)), 99 | ?assertEqual({ok, V1}, get(B, K)), 100 | 101 | %% put {K, V2} 102 | ?assertEqual(ok, put(B, K, V2)), 103 | ?assertEqual({ok, V2}, get(B, K)), 104 | 105 | %% verify K is in the list of keys 106 | ?assertEqual({ok, [K]}, list_keys(B)), 107 | 108 | %% delete K 109 | ?assertEqual(ok, delete(B, K)), 110 | ?assertMatch({error, notfound}, get(B, K)), 111 | 112 | %% stop the application 113 | application:stop(riakpool). 114 | 115 | -endif. 116 | -------------------------------------------------------------------------------- /src/riakpool_connection_sup.erl: -------------------------------------------------------------------------------- 1 | %% @author David Weldon 2 | %% @hidden 3 | 4 | -module(riakpool_connection_sup). 5 | -behaviour(supervisor). 6 | -export([start_link/0]). 7 | -export([init/1]). 8 | 9 | start_link() -> 10 | supervisor:start_link({local, ?MODULE}, ?MODULE, []). 11 | 12 | init([]) -> 13 | {ok, {{simple_one_for_one, 0, 1}, 14 | [{connections, {riakc_pb_socket, start_link, []}, 15 | temporary, brutal_kill, worker, [riakc_pb_socket]}]}}. 16 | -------------------------------------------------------------------------------- /src/riakpool_sup.erl: -------------------------------------------------------------------------------- 1 | %% @author David Weldon 2 | %% @hidden 3 | 4 | -module(riakpool_sup). 5 | -behaviour(supervisor). 6 | -export([start_link/0]). 7 | -export([init/1]). 8 | 9 | start_link() -> 10 | supervisor:start_link({local, ?MODULE}, ?MODULE, []). 11 | 12 | init([]) -> 13 | ConnectionSup = 14 | {riakpool_connection_sup, {riakpool_connection_sup, start_link, []}, 15 | permanent, 2000, supervisor, [riakpool_connection_sup]}, 16 | Pool = 17 | {riakpool, {riakpool, start_link, []}, 18 | permanent, 2000, worker, [riakpool]}, 19 | {ok, {{one_for_all, 5, 30}, [ConnectionSup, Pool]}}. 20 | --------------------------------------------------------------------------------