Skip to content

Add async listener initialization. #9

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
aef- wants to merge 1 commit into
base: openriak-3.4
Choose a base branch
from
Open

Add async listener initialization. #9

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
347 changes: 284 additions & 63 deletions src/riak_api_pb_listener.erl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,7 @@
%% -------------------------------------------------------------------
%%
%% riak_api_pb_listener: Listen for protocol buffer clients
%%
%% Copyright (c) 2007-2012 Basho Technologies, Inc. All Rights Reserved.
%% Copyright (c) 2007-2015 Basho Technologies, Inc.
%% Copyright (c) 2022-2023 Workday, Inc.
%%
%% This file is provided to you under the Apache License,
%% Version 2.0 (the "License"); you may not use this file
Expand All @@ -19,79 +18,188 @@
%% under the License.
%%
%% -------------------------------------------------------------------

%% @doc entry point for TCP-based protocol buffers service

%%
%% @doc Entry point for TCP-based protocol buffers service.
%%
%% This is a drop-in replacement for the upstream `riak_api_pb_listener'
%% module that incorporates its own non-blocking socket acceptance behavior
%% rather than relying on `gen_nb_server'. This approach allows us to simplify
%% the implementation while taking advantage of the continuation pattern added
%% to `gen_server' in OTP 21 that eliminates a potential race condition.
%%
%% Additionally, this implementation initializes asynchronously so that it can
%% wait for the KV service to be up before accepting connections, resulting in
%% clients getting consistent `{tcp, econnrefused}' errors instead of various
%% error results prior to the service being ready.
%%
%% @todo Consider for open-source submission, without the wd_ prefix.
%%
%% As currently implemented, a number of `riak_api_pb_listener' functions are
%% used in order to maintain compatibility should that module's behavior
%% change. In order to replace `riak_api_pb_listener' in place, those
%% functions' implementations and unit tests would first need to be moved into
%% this module.
%%
-module(riak_api_pb_listener).
-behaviour(gen_nb_server).
-export([start_link/2]).
-export([init/1, handle_call/3, handle_cast/2, handle_info/2,
terminate/2, code_change/3]).
-export([sock_opts/0, new_connection/2]).
-export([get_listeners/0]).
-behaviour(gen_server).

-include_lib("kernel/include/logger.hrl").

-record(state, {portnum}).
% Public API
-export([
start_link/2
]).

%% @doc Starts the PB listener
-spec start_link(inet:ip_address() | string(), non_neg_integer()) -> {ok, pid()} | {error, term()}.
start_link(IpAddr, PortNum) ->
gen_nb_server:start_link(?MODULE, IpAddr, PortNum, [PortNum]).
% gen_server callbacks
-export([
handle_continue/2,
handle_call/3,
handle_cast/2,
handle_info/2,
init/1,
terminate/2
]).

%% @doc Initialization callback for gen_nb_server.
-spec init(list()) -> {ok, #state{}}.
init([PortNum]) ->
{ok, #state{portnum=PortNum}}.
% riak_api_sup callback
-export([
get_listeners/0
]).

%% @doc Preferred socket options for the listener.
-spec sock_opts() -> [gen_tcp:option()].
sock_opts() ->
BackLog = app_helper:get_env(riak_api, pb_backlog, 128),
NoDelay = app_helper:get_env(riak_api, disable_pb_nagle, true),
KeepAlive = app_helper:get_env(riak_api, pb_keepalive, true),
[binary, {packet, raw}, {reuseaddr, true}, {backlog, BackLog}, {nodelay, NoDelay}, {keepalive, KeepAlive}].
%% ===================================================================
%% Types
%% ===================================================================

%% Use a unique name for the state record to avoid confusion if the original
%% module's functions are implemented here.
-record(wdpbl, {
sock :: gen_tcp:socket() | undefined,
ip :: inet:ip_address(),
port :: inet:port_number()
}).

-type pb_addr() :: inet:ip_address() | nonempty_string().
-type pb_port() :: inet:port_number().
-type state() :: #wdpbl{}.

%% Milliseconds between checks in wait_for_node_kv_service/1.
-define(WAIT_FOR_NODE_KV_RETRY_INTERVAL, 277).

%% @doc The handle_call/3 gen_nb_server callback. Unused.
-spec handle_call(term(), {pid(),_}, #state{}) -> {reply, term(), #state{}}.
handle_call(_Req, _From, State) ->
%% Milliseconds between checks for riak_core_node_watcher.
-define(WAIT_FOR_NODE_WATCHER_RETRY_INTERVAL, 181).

%% Message from init/1 to handle_continue/2
-define(ASYNC_INIT_CONTINUE_MSG, {?MODULE, init_tcp_server_with_wait}).

%% Simplified guards
-define(is_ip_addr(A), (erlang:is_tuple(A)
andalso (erlang:size(A) =:= 4 orelse erlang:size(A) =:= 8))).
-define(is_port_num(P),
(erlang:is_integer(P) andalso P >= 0 andalso P =< 65535)).

%% ===================================================================
%% Public API
%% ===================================================================

%% @doc Starts the PB listener with validated state.
%%
%% Throws a `badarg' error in the calling process if parameters are not valid.
-spec start_link(IpAddr :: pb_addr(), Port :: pb_port())
-> {ok, pid()} | {error, term()}.
start_link(IpAddr, Port) ->
State = init_state(IpAddr, Port),
gen_server:start_link(?MODULE, State, []).

%% ===================================================================
%% gen_server callbacks
%% ===================================================================

%% @doc Initialization callback for `gen_server' behavior.
%%
%% `State' has been validated by `start_link/2' before this is called.
-spec init(State :: state()) -> {ok, state(), {continue, term()}}.
init(State) ->
{ok, State, {continue, ?ASYNC_INIT_CONTINUE_MSG}}.

%% @doc Continuation callback for `gen_server' behavior.
%%
%% Performs asynchronous initialization of the listener.
-spec handle_continue(Continue :: term(), State :: state())
-> {noreply, state()} | {stop, term(), state()}.
handle_continue(?ASYNC_INIT_CONTINUE_MSG,
#wdpbl{ip = IpAddr, port = Port, sock = undefined} = State) ->
%% This call only returns on success, it blocks forever otherwise.
ok = wait_for_listener_ready(),
SockOpts = [{ip, IpAddr} | sock_opts()],
case gen_tcp:listen(Port, SockOpts) of
{ok, Socket} ->
case prim_inet:async_accept(Socket, -1) of
{ok, _Ref} ->
{noreply, State#wdpbl{sock = Socket}};
{error, AError} ->
_ = gen_tcp:close(Socket),
{stop, AError, State}
end;
{error, LError} ->
{stop, LError, State}
end;
handle_continue(Continue, State) ->
?LOG_ERROR("unhandled continuation ~0p", [Continue]),
{stop, {badarg, [Continue, State]}, State}.

%% @doc Unused required `gen_server' callback.
-spec handle_call(term(), {pid(), term()}, state())
-> {reply, term(), state()}.
handle_call(Request, From, State) ->
?LOG_WARNING("unhandled request ~0p from ~0p", [Request, From]),
{reply, not_implemented, State}.

%% @doc The handle_cast/2 gen_nb_server callback. Unused.
-spec handle_cast(term(), #state{}) -> {noreply, #state{}}.
handle_cast(_Msg, State) -> {noreply, State}.
%% @doc Unused required `gen_server' callback.
-spec handle_cast(Message :: term(), State :: state()) -> {noreply, state()}.
handle_cast(Message, State) ->
?LOG_WARNING("unhandled message ~0p", [Message]),
{noreply, State}.

%% @doc The handle_info/2 gen_nb_server callback. Unused.
-spec handle_info(term(), #state{}) -> {noreply, #state{}}.
handle_info(_Info, State) -> {noreply, State}.
%% @doc Message callback for `gen_server' behavior.
%%
%% Wires accepted socket connections through to their handlers.
-spec handle_info(Message :: term(), State :: state())
-> {noreply, state()} | {stop, term(), state()}.
handle_info({inet_async, ListSock, _Ref, {ok, CliSocket}}, StateIn) ->
case inet_db:register_socket(CliSocket, inet_tcp) of
true ->
{ok, StateOut} = new_connection(CliSocket, StateIn),
case prim_inet:async_accept(ListSock, -1) of
{ok, _} ->
{noreply, StateOut};
{error, Reason} ->
{stop, Reason, StateOut}
end;
_ ->
?LOG_ERROR("Failed to register socket ~w", [CliSocket]),
_ = gen_tcp:close(CliSocket),
{stop, {badarg, [CliSocket]}, StateIn}
end;
handle_info(Message, State) ->
?LOG_WARNING("Unhandled message ~0p", [Message]),
{noreply, State}.

%% @doc The code_change/3 gen_nb_server callback. Unused.
-spec terminate(Reason, State) -> ok when
Reason :: normal | shutdown | {shutdown,term()} | term(),
State :: #state{}.
terminate(_Reason, _State) ->
ok.
%% @doc Termination callback for `gen_server' behavior.
-spec terminate(Reason :: term(), State :: state()) -> Ignored :: term().
terminate(Reason, #wdpbl{sock = undefined} = State) ->
%% If the socket is undefined then we're in the asynchronous
%% initialization phase and handle_continue/2 hasn't been called yet,
%% so there's nothing to clean up.
?LOG_DEBUG("terminate(~0p) with state: ~0p", [Reason, State]);
terminate(Reason, #wdpbl{sock = Socket} = State) ->
_ = gen_tcp:close(Socket),
?LOG_DEBUG("terminate(~0p) with state: ~0p", [Reason, State]).

%% @doc The gen_server code_change/3 callback, called when performing
%% a hot code upgrade on the server. Currently unused.
-spec code_change(OldVsn, State, Extra) -> {ok, State} | {error, Reason}
when
OldVsn :: Vsn | {down, Vsn},
Vsn :: term(),
State :: #state{},
Extra :: term(),
Reason :: term().
code_change(_OldVsn, State, _Extra) -> {ok, State}.

%% @doc The connection initiation callback for gen_nb_server, called
%% when a new socket is accepted.
-spec new_connection(port(), #state{}) -> {ok, #state{}}.
new_connection(Socket, State) ->
{ok, Pid} = riak_api_pb_sup:start_socket(),
ok = gen_tcp:controlling_process(Socket, Pid),
ok = riak_api_pb_server:set_socket(Pid, Socket),
{ok, State}.
%% ===================================================================
%% riak_api_sup callback
%% ===================================================================

%% @doc Returns the endpoint upon which this service will be initialized.
-spec get_listeners() -> list({inet:ip_address(), inet:port_number()}).
get_listeners() ->
DefaultListener = case {get_ip(), get_port()} of
{undefined, _} -> [];
Expand Down Expand Up @@ -125,9 +233,122 @@ get_ip() ->
IP
end.

%% ===================================================================
%% Wait for the KV service to be ready
%% ===================================================================

%% @hidden
%% Wait until required services are online, or forever if they don't start.
%%
%% The riak_api EUnit tests don't start riak_core/riak_kv, so this would
%% block forever under TEST, preventing the listener from ever opening its
%% socket.
-spec wait_for_listener_ready() -> ok.
-ifdef(TEST).
wait_for_listener_ready() ->
ok.
-else. % ! TEST
wait_for_listener_ready() ->
wait_for_node_service_watcher(),
wait_for_node_kv_service(erlang:node()).
-endif. % TEST

-ifndef(TEST).
%% @hidden
%% Returns `ok' only when running services can be checked safely.
%%
%% This is waaay too tightly coupled to the riak_core_node_watcher
%% implementation, but there doesn't seem to be any way to avoid that.
%%
%% As currently implemented, riak_core_node_watcher:services/1 relies on the
%% ets table directly, not on the running gen_server.
%% We'd prefer to just check for the running gen_server by name, since its
%% init/1 function initializes the ets table, but gen_server registers the
%% name *before* running init/1, so erlang:whereis can return a pid before
%% the ets table has actually been initialized.
%%
%% We use the services/0 function, which performs a lightweight query on the
%% ets table via a gen_server call, as a proxy for determining that services/1
%% can be invoked, so that if the services/1 implementation changes to go
%% through the gen_server this strategy should still be safe on the assumption
%% that services/1 and services/0 would almost certainly operate from the same
%% shared state (as they do now).
%%
%% If riak_core_node_watcher:init/1 fails for any reason that will cascade to
%% this process crashing as well, and both gen_servers will be restarted by
%% their respective supervisors to try it all again.
%%
-spec wait_for_node_service_watcher() -> ok.
wait_for_node_service_watcher() ->
RegisteredService = riak_core_node_watcher,
?LOG_DEBUG("checking ~s", [RegisteredService]),
case erlang:whereis(RegisteredService) of
undefined ->
timer:sleep(?WAIT_FOR_NODE_WATCHER_RETRY_INTERVAL),
wait_for_node_service_watcher();
_ ->
%% We don't care about the result, this is just the cheapest way
%% to wait for the gen_server to finish its (and the ets table's)
%% initialization.
_ = riak_core_node_watcher:services(),
ok
end.

%% @hidden
%% Returns `ok' only when listener requests can be handled successfully
-spec wait_for_node_kv_service(Node :: node()) -> ok.
wait_for_node_kv_service(Node) ->
?LOG_DEBUG("checking for KV service on ~w", [Node]),
case lists:member(riak_kv, riak_core_node_watcher:services(Node)) of
true ->
ok;
_ ->
timer:sleep(?WAIT_FOR_NODE_KV_RETRY_INTERVAL),
wait_for_node_kv_service(Node)
end.
-endif. % ! TEST

%% ===================================================================
%% Internal functions
%% ===================================================================

-spec init_state(IpAddr :: pb_addr(), Port :: pb_port())
-> state() | no_return().
%% @private
%% Initializes the state record with validated parameters, allowing the
%% `start' function(s) to throw a `badarg' error *before* spawning the gs
%% process, yielding an informative stack trace.
init_state(_IpAddr, Port) when not ?is_port_num(Port) ->
erlang:error(badarg, [Port]);
init_state(IpAddrStr, Port) when not ?is_ip_addr(IpAddrStr) ->
case inet:parse_address(IpAddrStr) of
{ok, IpAddr} ->
init_state(IpAddr, Port);
_ ->
erlang:error(badarg, [IpAddrStr])
end;
init_state(IpAddr, Port) ->
#wdpbl{ip = IpAddr, port = Port}.

%% @private Called when a new socket is accepted.
-spec new_connection(Socket :: port(), state()) -> {ok, state()}.
new_connection(Socket, State) ->
{ok, Pid} = riak_api_pb_sup:start_socket(),
ok = gen_tcp:controlling_process(Socket, Pid),
ok = riak_api_pb_server:set_socket(Pid, Socket),
{ok, State}.

%% @private Preferred socket options for the listener.
-spec sock_opts() -> list(gen_tcp:listen_option()).
sock_opts() ->
BackLog = app_helper:get_env(riak_api, pb_backlog, 128),
NoDelay = app_helper:get_env(riak_api, disable_pb_nagle, true),
KeepAlive = app_helper:get_env(riak_api, pb_keepalive, true),
[binary, {packet, raw}, {reuseaddr, true}, {backlog, BackLog},
{nodelay, NoDelay}, {keepalive, KeepAlive}].

-ifdef(TEST).
-include_lib("eunit/include/eunit.hrl").
-compile([export_all, nowarn_export_all]).

listeners_test_() ->
{foreach,
Expand Down
Loading
Loading