src/rexi/src/rexi.erl - couchdb - Git at Google

 % Licensed under the Apache License, Version 2.0 (the "License"); you may not
 % use this file except in compliance with the License. You may obtain a copy of
 % the License at
 %
 %   http://www.apache.org/licenses/LICENSE-2.0
 %
 % Unless required by applicable law or agreed to in writing, software
 % distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 % WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 % License for the specific language governing permissions and limitations under
 % the License.

 -module(rexi).
 -export([start/0, stop/0, restart/0]).
 -export([cast/2, cast/3, cast/4, kill/2, kill_all/1]).
 -export([reply/1, sync_reply/1, sync_reply/2]).
 -export([async_server_call/2, async_server_call/3]).
 -export([stream_init/0, stream_init/1]).
 -export([stream_start/1, stream_cancel/1]).
 -export([stream/1, stream/2, stream/3, stream_ack/1, stream_ack/2]).
 -export([stream2/1, stream2/2, stream2/3, stream_last/1, stream_last/2]).
 -export([ping/0]).

 -include_lib("rexi/include/rexi.hrl").

 start() ->
     application:start(rexi).

 stop() ->
     application:stop(rexi).

 restart() ->
     stop(), start().


 %% @equiv cast(Node, self(), MFA)
 -spec cast(node(), {atom(), atom(), list()}) -> reference().
 cast(Node, MFA) ->
     cast(Node, self(), MFA).

 %% @doc Executes apply(M, F, A) on Node.
 %% You might want to use this instead of rpc:cast/4 for two reasons.  First,
 %% the Caller pid and the returned reference are inserted into the remote
 %% process' dictionary as `rexi_from', so it has a way to communicate with you.
 %% Second, the remote process is monitored. If it exits with a Reason other
 %% than normal, Caller will receive a message of the form
 %% `{Ref, {rexi_EXIT, Reason}}' where Ref is the returned reference.
 -spec cast(node(), pid(), {atom(), atom(), list()}) -> reference().
 cast(Node, Caller, MFA) ->
     Ref = make_ref(),
     Msg = cast_msg({doit, {Caller, Ref}, get(nonce), MFA}),
     rexi_utils:send(rexi_utils:server_pid(Node), Msg),
     Ref.

 %% @doc Executes apply(M, F, A) on Node.
 %% This version accepts a sync option which uses the erlang:send/2 call
 %% directly in process instead of deferring to a spawned process if
 %% erlang:send/2 were to block. If the sync option is omitted this call
 %% is identical to cast/3.
 -spec cast(node(), pid(), {atom(), atom(), list()}, [atom()]) -> reference().
 cast(Node, Caller, MFA, Options) ->
     case lists:member(sync, Options) of
         true ->
             Ref = make_ref(),
             Msg = cast_msg({doit, {Caller, Ref}, get(nonce), MFA}),
             erlang:send(rexi_utils:server_pid(Node), Msg),
             Ref;
         false ->
             cast(Node, Caller, MFA)
     end.

 %% @doc Sends an async kill signal to the remote process associated with Ref.
 %% No rexi_EXIT message will be sent.
 -spec kill(node(), reference()) -> ok.
 kill(Node, Ref) ->
     rexi_utils:send(rexi_utils:server_pid(Node), cast_msg({kill, Ref})),
     ok.

 %% @doc Sends an async kill signal to the remote processes associated with Refs.
 %% No rexi_EXIT message will be sent.
 -spec kill_all([{node(), reference()}]) -> ok.
 kill_all(NodeRefs) when is_list(NodeRefs) ->
     %% Upgrade clause. Since kill_all is a new message, nodes in a mixed
     %% cluster won't know how to process it. In that case, the default is to send
     %% the individual kill messages. Once all the nodes have been upgraded, can
     %% configure the cluster to send kill_all messages.
     case config:get_boolean("rexi", "use_kill_all", false) of
         true ->
             PerNodeMap = lists:foldl(fun({Node, Ref}, Acc) ->
                 maps:update_with(Node, fun(Refs) ->
                     [Ref | Refs]
                 end, [Ref], Acc)
             end, #{}, NodeRefs),
             maps:map(fun(Node, Refs) ->
                 ServerPid = rexi_utils:server_pid(Node),
                 rexi_utils:send(ServerPid, cast_msg({kill_all, Refs}))
             end, PerNodeMap);
         false ->
             lists:foreach(fun({Node, Ref}) -> kill(Node, Ref) end, NodeRefs)
     end,
     ok.

 %% @equiv async_server_call(Server, self(), Request)
 -spec async_server_call(pid() | {atom(),node()}, any()) -> reference().
 async_server_call(Server, Request) ->
     async_server_call(Server, self(), Request).

 %% @doc Sends a properly formatted gen_server:call Request to the Server and
 %% returns the reference which the Server will include in its reply.  The
 %% function acts more like cast() than call() in that the server process
 %% is not monitored.  Clients who want to know if the server is alive should
 %% monitor it themselves before calling this function.
 -spec async_server_call(pid() | {atom(),node()}, pid(), any()) -> reference().
 async_server_call(Server, Caller, Request) ->
     Ref = make_ref(),
     rexi_utils:send(Server, {'$gen_call', {Caller,Ref}, Request}),
     Ref.

 %% @doc convenience function to reply to the original rexi Caller.
 -spec reply(any()) -> any().
 reply(Reply) ->
     {Caller, Ref} = get(rexi_from),
     erlang:send(Caller, {Ref,Reply}).

 %% @equiv sync_reply(Reply, 300000)
 sync_reply(Reply) ->
     sync_reply(Reply, 300000).

 %% @doc convenience function to reply to caller and wait for response.  Message
 %% is of the form {OriginalRef, {self(),reference()}, Reply}, which enables the
 %% original caller to respond back.
 -spec sync_reply(any(), pos_integer() | infinity) -> any().
 sync_reply(Reply, Timeout) ->
     {Caller, Ref} = get(rexi_from),
     Tag = make_ref(),
     erlang:send(Caller, {Ref, {self(),Tag}, Reply}),
     receive {Tag, Response} ->
         Response
     after Timeout ->
         timeout
     end.

 %% @equiv stream_init(300000)
 stream_init() ->
     stream_init(300000).

 %% @doc Initialize an RPC stream that involves sending multiple
 %% messages back to the coordinator.
 %%
 %% This should be called by rexi workers. It blocks until the
 %% coordinator responds with whether this worker should proceed.
 %% This function will either return with `ok` or call
 %% `erlang:exit/1`.
 -spec stream_init(pos_integer()) -> ok.
 stream_init(Timeout) ->
     case sync_reply(rexi_STREAM_INIT, Timeout) of
         rexi_STREAM_START ->
             ok;
         rexi_STREAM_CANCEL ->
             exit(normal);
         timeout ->
             couch_stats:increment_counter(
                 [rexi, streams, timeout, init_stream]
             ),
             exit(timeout);
         Else ->
             exit({invalid_stream_message, Else})
     end.

 %% @doc Start a worker stream
 %%
 %% If a coordinator wants to continue using a streaming worker it
 %% should use this function to inform the worker to continue
 %% sending messages. The `From` should be the value provided by
 %% the worker in the rexi_STREAM_INIT message.
 -spec stream_start({pid(), any()}) -> ok.
 stream_start({Pid, _Tag}=From) when is_pid(Pid) ->
     gen_server:reply(From, rexi_STREAM_START).

 %% @doc Cancel a worker stream
 %%
 %% If a coordinator decideds that a worker is not going to be part
 %% of the response it should use this function to cancel the worker.
 %% The `From` should be the value provided by the worker in the
 %% rexi_STREAM_INIT message.
 -spec stream_cancel({pid(), any()}) -> ok.
 stream_cancel({Pid, _Tag}=From) when is_pid(Pid) ->
     gen_server:reply(From, rexi_STREAM_CANCEL).

 %% @equiv stream(Msg, 100, 300000)
 stream(Msg) ->
     stream(Msg, 10, 300000).

 %% @equiv stream(Msg, Limit, 300000)
 stream(Msg, Limit) ->
     stream(Msg, Limit, 300000).

 %% @doc convenience function to stream messages to caller while blocking when
 %% a specific number of messages are outstanding. Message is of the form
 %% {OriginalRef, self(), Reply}, which enables the original caller to ack.
 -spec stream(any(), integer(), pos_integer() | infinity) -> any().
 stream(Msg, Limit, Timeout) ->
     try maybe_wait(Limit, Timeout) of
         {ok, Count} ->
             put(rexi_unacked, Count+1),
             {Caller, Ref} = get(rexi_from),
             erlang:send(Caller, {Ref, self(), Msg}),
             ok
     catch throw:timeout ->
         couch_stats:increment_counter([rexi, streams, timeout, stream]),
         exit(timeout)
     end.

 %% @equiv stream2(Msg, 5, 300000)
 stream2(Msg) ->
     Limit = config:get_integer("rexi", "stream_limit", 5),
     stream2(Msg, Limit).

 %% @equiv stream2(Msg, Limit, 300000)
 stream2(Msg, Limit) ->
     stream2(Msg, Limit, 300000).

 %% @doc Stream a message back to the coordinator. It limits the
 %% number of unacked messsages to Limit and throws a timeout error
 %% if it doesn't receive an ack in Timeout milliseconds. This
 %% is a combination of the old stream_start and stream functions
 %% which automatically does the stream initialization logic.
 -spec stream2(any(), pos_integer(), pos_integer() | inifinity) -> any().
 stream2(Msg, Limit, Timeout) ->
     maybe_init_stream(Timeout),
     try maybe_wait(Limit, Timeout) of
         {ok, Count} ->
             put(rexi_unacked, Count+1),
             {Caller, Ref} = get(rexi_from),
             erlang:send(Caller, {Ref, self(), Msg}),
             ok
     catch throw:timeout ->
         couch_stats:increment_counter([rexi, streams, timeout, stream]),
         exit(timeout)
     end.

 %% @equiv stream_last(Msg, 300000)
 stream_last(Msg) ->
     stream_last(Msg, 300000).

 %% @doc Send the last message in a stream. This difference between
 %% this and stream is that it uses rexi:reply/1 which doesn't include
 %% the worker pid and doesn't wait for a response from the controller.
 stream_last(Msg, Timeout) ->
     maybe_init_stream(Timeout),
     rexi:reply(Msg),
     ok.

 %% @equiv stream_ack(Client, 1)
 stream_ack(Client) ->
     erlang:send(Client, {rexi_ack, 1}).

 %% @doc Ack streamed messages
 stream_ack(Client, N) ->
     erlang:send(Client, {rexi_ack, N}).


 %% Sends a ping message to the coordinator. This is for long running
 %% operations on a node that could exceed the rexi timeout
 ping() ->
     {Caller, _} = get(rexi_from),
     erlang:send(Caller, {rexi, '$rexi_ping'}).


 %% internal functions %%

 cast_msg(Msg) -> {'$gen_cast', Msg}.

 maybe_init_stream(Timeout) ->
     case get(rexi_STREAM_INITED) of
         true ->
             ok;
         _ ->
             init_stream(Timeout)
     end.

 init_stream(Timeout) ->
     case sync_reply(rexi_STREAM_INIT, Timeout) of
         rexi_STREAM_START ->
             put(rexi_STREAM_INITED, true),
             ok;
         rexi_STREAM_CANCEL ->
             exit(normal);
         timeout ->
             exit(timeout);
         Else ->
             exit({invalid_stream_message, Else})
     end.

 maybe_wait(Limit, Timeout) ->
     case get(rexi_unacked) of
         undefined ->
             {ok, 0};
         Count when Count >= Limit ->
             wait_for_ack(Count, Timeout);
         Count ->
             drain_acks(Count)
     end.

 wait_for_ack(Count, Timeout) ->
     receive
         {rexi_ack, N} -> drain_acks(Count-N)
     after Timeout ->
         couch_stats:increment_counter([rexi, streams, timeout, wait_for_ack]),
         throw(timeout)
     end.

 drain_acks(Count) when Count < 0 ->
     erlang:error(mismatched_rexi_ack);
 drain_acks(Count) ->
     receive
         {rexi_ack, N} -> drain_acks(Count-N)
     after 0 ->
         {ok, Count}
     end.
	% Licensed under the Apache License, Version 2.0 (the "License"); you may not
	% use this file except in compliance with the License. You may obtain a copy of
	% the License at
	%
	% http://www.apache.org/licenses/LICENSE-2.0
	%
	% Unless required by applicable law or agreed to in writing, software
	% distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
	% WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
	% License for the specific language governing permissions and limitations under
	% the License.

	-module(rexi).
	-export([start/0, stop/0, restart/0]).
	-export([cast/2, cast/3, cast/4, kill/2, kill_all/1]).
	-export([reply/1, sync_reply/1, sync_reply/2]).
	-export([async_server_call/2, async_server_call/3]).
	-export([stream_init/0, stream_init/1]).
	-export([stream_start/1, stream_cancel/1]).
	-export([stream/1, stream/2, stream/3, stream_ack/1, stream_ack/2]).
	-export([stream2/1, stream2/2, stream2/3, stream_last/1, stream_last/2]).
	-export([ping/0]).

	-include_lib("rexi/include/rexi.hrl").

	start() ->
	application:start(rexi).

	stop() ->
	application:stop(rexi).

	restart() ->
	stop(), start().


	%% @equiv cast(Node, self(), MFA)
	-spec cast(node(), {atom(), atom(), list()}) -> reference().
	cast(Node, MFA) ->
	cast(Node, self(), MFA).

	%% @doc Executes apply(M, F, A) on Node.
	%% You might want to use this instead of rpc:cast/4 for two reasons. First,
	%% the Caller pid and the returned reference are inserted into the remote
	%% process' dictionary as `rexi_from', so it has a way to communicate with you.
	%% Second, the remote process is monitored. If it exits with a Reason other
	%% than normal, Caller will receive a message of the form
	%% `{Ref, {rexi_EXIT, Reason}}' where Ref is the returned reference.
	-spec cast(node(), pid(), {atom(), atom(), list()}) -> reference().
	cast(Node, Caller, MFA) ->
	Ref = make_ref(),
	Msg = cast_msg({doit, {Caller, Ref}, get(nonce), MFA}),
	rexi_utils:send(rexi_utils:server_pid(Node), Msg),
	Ref.

	%% @doc Executes apply(M, F, A) on Node.
	%% This version accepts a sync option which uses the erlang:send/2 call
	%% directly in process instead of deferring to a spawned process if
	%% erlang:send/2 were to block. If the sync option is omitted this call
	%% is identical to cast/3.
	-spec cast(node(), pid(), {atom(), atom(), list()}, [atom()]) -> reference().
	cast(Node, Caller, MFA, Options) ->
	case lists:member(sync, Options) of
	true ->
	Ref = make_ref(),
	Msg = cast_msg({doit, {Caller, Ref}, get(nonce), MFA}),
	erlang:send(rexi_utils:server_pid(Node), Msg),
	Ref;
	false ->
	cast(Node, Caller, MFA)
	end.

	%% @doc Sends an async kill signal to the remote process associated with Ref.
	%% No rexi_EXIT message will be sent.
	-spec kill(node(), reference()) -> ok.
	kill(Node, Ref) ->
	rexi_utils:send(rexi_utils:server_pid(Node), cast_msg({kill, Ref})),
	ok.

	%% @doc Sends an async kill signal to the remote processes associated with Refs.
	%% No rexi_EXIT message will be sent.
	-spec kill_all([{node(), reference()}]) -> ok.
	kill_all(NodeRefs) when is_list(NodeRefs) ->
	%% Upgrade clause. Since kill_all is a new message, nodes in a mixed
	%% cluster won't know how to process it. In that case, the default is to send
	%% the individual kill messages. Once all the nodes have been upgraded, can
	%% configure the cluster to send kill_all messages.
	case config:get_boolean("rexi", "use_kill_all", false) of
	true ->
	PerNodeMap = lists:foldl(fun({Node, Ref}, Acc) ->
	maps:update_with(Node, fun(Refs) ->
	[Ref \| Refs]
	end, [Ref], Acc)
	end, #{}, NodeRefs),
	maps:map(fun(Node, Refs) ->
	ServerPid = rexi_utils:server_pid(Node),
	rexi_utils:send(ServerPid, cast_msg({kill_all, Refs}))
	end, PerNodeMap);
	false ->
	lists:foreach(fun({Node, Ref}) -> kill(Node, Ref) end, NodeRefs)
	end,
	ok.

	%% @equiv async_server_call(Server, self(), Request)
	-spec async_server_call(pid() \| {atom(),node()}, any()) -> reference().
	async_server_call(Server, Request) ->
	async_server_call(Server, self(), Request).

	%% @doc Sends a properly formatted gen_server:call Request to the Server and
	%% returns the reference which the Server will include in its reply. The
	%% function acts more like cast() than call() in that the server process
	%% is not monitored. Clients who want to know if the server is alive should
	%% monitor it themselves before calling this function.
	-spec async_server_call(pid() \| {atom(),node()}, pid(), any()) -> reference().
	async_server_call(Server, Caller, Request) ->
	Ref = make_ref(),
	rexi_utils:send(Server, {'$gen_call', {Caller,Ref}, Request}),
	Ref.

	%% @doc convenience function to reply to the original rexi Caller.
	-spec reply(any()) -> any().
	reply(Reply) ->
	{Caller, Ref} = get(rexi_from),
	erlang:send(Caller, {Ref,Reply}).

	%% @equiv sync_reply(Reply, 300000)
	sync_reply(Reply) ->
	sync_reply(Reply, 300000).

	%% @doc convenience function to reply to caller and wait for response. Message
	%% is of the form {OriginalRef, {self(),reference()}, Reply}, which enables the
	%% original caller to respond back.
	-spec sync_reply(any(), pos_integer() \| infinity) -> any().
	sync_reply(Reply, Timeout) ->
	{Caller, Ref} = get(rexi_from),
	Tag = make_ref(),
	erlang:send(Caller, {Ref, {self(),Tag}, Reply}),
	receive {Tag, Response} ->
	Response
	after Timeout ->
	timeout
	end.

	%% @equiv stream_init(300000)
	stream_init() ->
	stream_init(300000).

	%% @doc Initialize an RPC stream that involves sending multiple
	%% messages back to the coordinator.
	%%
	%% This should be called by rexi workers. It blocks until the
	%% coordinator responds with whether this worker should proceed.
	%% This function will either return with `ok` or call
	%% `erlang:exit/1`.
	-spec stream_init(pos_integer()) -> ok.
	stream_init(Timeout) ->
	case sync_reply(rexi_STREAM_INIT, Timeout) of
	rexi_STREAM_START ->
	ok;
	rexi_STREAM_CANCEL ->
	exit(normal);
	timeout ->
	couch_stats:increment_counter(
	[rexi, streams, timeout, init_stream]
	),
	exit(timeout);
	Else ->
	exit({invalid_stream_message, Else})
	end.

	%% @doc Start a worker stream
	%%
	%% If a coordinator wants to continue using a streaming worker it
	%% should use this function to inform the worker to continue
	%% sending messages. The `From` should be the value provided by
	%% the worker in the rexi_STREAM_INIT message.
	-spec stream_start({pid(), any()}) -> ok.
	stream_start({Pid, _Tag}=From) when is_pid(Pid) ->
	gen_server:reply(From, rexi_STREAM_START).

	%% @doc Cancel a worker stream
	%%
	%% If a coordinator decideds that a worker is not going to be part
	%% of the response it should use this function to cancel the worker.
	%% The `From` should be the value provided by the worker in the
	%% rexi_STREAM_INIT message.
	-spec stream_cancel({pid(), any()}) -> ok.
	stream_cancel({Pid, _Tag}=From) when is_pid(Pid) ->
	gen_server:reply(From, rexi_STREAM_CANCEL).

	%% @equiv stream(Msg, 100, 300000)
	stream(Msg) ->
	stream(Msg, 10, 300000).

	%% @equiv stream(Msg, Limit, 300000)
	stream(Msg, Limit) ->
	stream(Msg, Limit, 300000).

	%% @doc convenience function to stream messages to caller while blocking when
	%% a specific number of messages are outstanding. Message is of the form
	%% {OriginalRef, self(), Reply}, which enables the original caller to ack.
	-spec stream(any(), integer(), pos_integer() \| infinity) -> any().
	stream(Msg, Limit, Timeout) ->
	try maybe_wait(Limit, Timeout) of
	{ok, Count} ->
	put(rexi_unacked, Count+1),
	{Caller, Ref} = get(rexi_from),
	erlang:send(Caller, {Ref, self(), Msg}),
	ok
	catch throw:timeout ->
	couch_stats:increment_counter([rexi, streams, timeout, stream]),
	exit(timeout)
	end.

	%% @equiv stream2(Msg, 5, 300000)
	stream2(Msg) ->
	Limit = config:get_integer("rexi", "stream_limit", 5),
	stream2(Msg, Limit).

	%% @equiv stream2(Msg, Limit, 300000)
	stream2(Msg, Limit) ->
	stream2(Msg, Limit, 300000).

	%% @doc Stream a message back to the coordinator. It limits the
	%% number of unacked messsages to Limit and throws a timeout error
	%% if it doesn't receive an ack in Timeout milliseconds. This
	%% is a combination of the old stream_start and stream functions
	%% which automatically does the stream initialization logic.
	-spec stream2(any(), pos_integer(), pos_integer() \| inifinity) -> any().
	stream2(Msg, Limit, Timeout) ->
	maybe_init_stream(Timeout),
	try maybe_wait(Limit, Timeout) of
	{ok, Count} ->
	put(rexi_unacked, Count+1),
	{Caller, Ref} = get(rexi_from),
	erlang:send(Caller, {Ref, self(), Msg}),
	ok
	catch throw:timeout ->
	couch_stats:increment_counter([rexi, streams, timeout, stream]),
	exit(timeout)
	end.

	%% @equiv stream_last(Msg, 300000)
	stream_last(Msg) ->
	stream_last(Msg, 300000).

	%% @doc Send the last message in a stream. This difference between
	%% this and stream is that it uses rexi:reply/1 which doesn't include
	%% the worker pid and doesn't wait for a response from the controller.
	stream_last(Msg, Timeout) ->
	maybe_init_stream(Timeout),
	rexi:reply(Msg),
	ok.

	%% @equiv stream_ack(Client, 1)
	stream_ack(Client) ->
	erlang:send(Client, {rexi_ack, 1}).

	%% @doc Ack streamed messages
	stream_ack(Client, N) ->
	erlang:send(Client, {rexi_ack, N}).


	%% Sends a ping message to the coordinator. This is for long running
	%% operations on a node that could exceed the rexi timeout
	ping() ->
	{Caller, _} = get(rexi_from),
	erlang:send(Caller, {rexi, '$rexi_ping'}).


	%% internal functions %%

	cast_msg(Msg) -> {'$gen_cast', Msg}.

	maybe_init_stream(Timeout) ->
	case get(rexi_STREAM_INITED) of
	true ->
	ok;
	_ ->
	init_stream(Timeout)
	end.

	init_stream(Timeout) ->
	case sync_reply(rexi_STREAM_INIT, Timeout) of
	rexi_STREAM_START ->
	put(rexi_STREAM_INITED, true),
	ok;
	rexi_STREAM_CANCEL ->
	exit(normal);
	timeout ->
	exit(timeout);
	Else ->
	exit({invalid_stream_message, Else})
	end.

	maybe_wait(Limit, Timeout) ->
	case get(rexi_unacked) of
	undefined ->
	{ok, 0};
	Count when Count >= Limit ->
	wait_for_ack(Count, Timeout);
	Count ->
	drain_acks(Count)
	end.

	wait_for_ack(Count, Timeout) ->
	receive
	{rexi_ack, N} -> drain_acks(Count-N)
	after Timeout ->
	couch_stats:increment_counter([rexi, streams, timeout, wait_for_ack]),
	throw(timeout)
	end.

	drain_acks(Count) when Count < 0 ->
	erlang:error(mismatched_rexi_ack);
	drain_acks(Count) ->
	receive
	{rexi_ack, N} -> drain_acks(Count-N)
	after 0 ->
	{ok, Count}
	end.