diff --git a/src/riak_api_pb_listener.erl b/src/riak_api_pb_listener.erl index 27fdcf5..b4adb99 100644 --- a/src/riak_api_pb_listener.erl +++ b/src/riak_api_pb_listener.erl @@ -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 @@ -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, _} -> []; @@ -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, diff --git a/src/riak_api_sup.erl b/src/riak_api_sup.erl index 5c2cd2f..61bc3c0 100644 --- a/src/riak_api_sup.erl +++ b/src/riak_api_sup.erl @@ -2,7 +2,8 @@ %% %% riak_api_sup: supervise the Riak API services %% -%% Copyright (c) 2007-2010 Basho Technologies, Inc. All Rights Reserved. +%% Copyright (c) 2012-2013 Basho Technologies, Inc. +%% Copyright (c) 2023 Workday, Inc. %% %% This file is provided to you under the Apache License, %% Version 2.0 (the "License"); you may not use this file @@ -34,9 +35,9 @@ -define(CHILD(I, Type), {I, {I, start_link, []}, permanent, 5000, Type, [I]}). -define(CHILD(I, Type, Args), {I, {I, start_link, Args}, permanent, 5000, Type, [I]}). -define(LNAME(IP, Port), lists:flatten(io_lib:format("pb://~p:~p", [IP, Port]))). --define(PB_LISTENER(IP, Port), {?LNAME(IP, Port), - {riak_api_pb_listener, start_link, [IP, Port]}, - permanent, 5000, worker, [riak_api_pb_listener]}). +-define(PB_LISTENER(Mod, IP, Port), {?LNAME(IP, Port), + {Mod, start_link, [IP, Port]}, + permanent, 5000, worker, [Mod]}). %% @doc Starts the supervisor. -spec start_link() -> {ok, pid()} | {error, term()}. start_link() -> @@ -51,7 +52,7 @@ start_link() -> init([]) -> Helper = ?CHILD(riak_api_pb_registration_helper, worker), Registrar = ?CHILD(riak_api_pb_registrar, worker), - PBProcesses = pb_processes(riak_api_pb_listener:get_listeners()), + PBProcesses = pb_processes(), WebProcesses = web_processes(riak_api_web:get_listeners()), NetworkProcesses = PBProcesses ++ WebProcesses, {ok, {{one_for_one, 10, 10}, [Helper, Registrar|NetworkProcesses]}}. @@ -70,11 +71,16 @@ web_listener_spec(Scheme, Binding) -> %% Generates child specs from the PB listener configuration. %% @private -pb_processes([]) -> - ?LOG_INFO("No PB listeners were configured, PB connections will be disabled."), - []; -pb_processes(Listeners) -> - [?CHILD(riak_api_pb_sup, supervisor)| pb_listener_specs(Listeners)]. - -pb_listener_specs(Pairs) -> - [ ?PB_LISTENER(IP, Port) || {IP, Port} <- Pairs ]. +-spec pb_processes() -> list(supervisor:child_spec()). +pb_processes() -> + case riak_api_pb_listener:get_listeners() of + [] -> + ?LOG_INFO( + "No PB listeners were configured," + " PB connections will be disabled."), + []; + Endpoints -> + ListenerSpecs = + [?PB_LISTENER(riak_api_pb_listener, IP, Port) || {IP, Port} <- Endpoints], + [?CHILD(riak_api_pb_sup, supervisor) | ListenerSpecs] + end.