src/fabric/src/fabric_streams.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(fabric_streams).

 -export([
     start/2,
     start/3,
     start/4,
     start/5,
     cleanup/1,
     enable_watchdog/0,
     kick_watchdog/0
 ]).

 -include_lib("fabric/include/fabric.hrl").
 -include_lib("mem3/include/mem3.hrl").

 -define(WORKER_WATCHDOG, fabric_worker_watchdog).
 -define(ADD_WORKER, add_worker).
 -define(WATCHDOG_ENABLE, watchdog_enable).
 -define(WATCHDOG_TIMEOUT, watchdog_timeout).
 -define(WATCHDOG_KICK, watchdog_kick).
 -define(WATCHDOG_LAST_KICK, watchdog_last_kick).

 -define(WATCHDOG_WAS_KICKED, St#watchdog_state.kicked).
 -define(WATCHDOG_IS_ENABLED, St#watchdog_state.timer /= undefined).

 -record(watchdog_state, {
     coordinator,
     workers,
     timer,
     kicked = false
 }).

 start(Workers, Keypos) ->
     start(Workers, Keypos, undefined, undefined).

 start(Workers, Keypos, RingOpts) ->
     start(Workers, Keypos, undefined, undefined, RingOpts).

 start(Workers, Keypos, StartFun, Replacements) ->
     start(Workers, Keypos, StartFun, Replacements, []).

 start(Workers0, Keypos, StartFun, Replacements, RingOpts) ->
     Fun = fun handle_stream_start/3,
     Acc = #stream_acc{
         workers = fabric_dict:init(Workers0, waiting),
         ready = [],
         start_fun = StartFun,
         replacements = Replacements,
         ring_opts = RingOpts
     },
     spawn_worker_watchdog(self(), Workers0),
     Timeout = fabric_util:request_timeout(),
     case rexi_utils:recv(Workers0, Keypos, Fun, Acc, Timeout, infinity) of
         {ok, #stream_acc{ready = Workers}} ->
             AckedWorkers = fabric_dict:fold(
                 fun(Worker, From, WorkerAcc) ->
                     rexi:stream_start(From),
                     [Worker | WorkerAcc]
                 end,
                 [],
                 Workers
             ),
             {ok, AckedWorkers};
         Else ->
             Else
     end.

 cleanup(Workers) ->
     % Stop the auxiliary watchdog process as we got to the point where cleanup
     % happens in the regular fashion so we don't want to send 2x the number kill
     % messages
     case get(?WORKER_WATCHDOG) of
         WatchdogPid when is_pid(WatchdogPid) ->
             erase(?WORKER_WATCHDOG),
             exit(WatchdogPid, kill);
         _ ->
             ok
     end,
     fabric_util:cleanup(Workers).

 handle_stream_start({rexi_DOWN, _, {_, NodeRef}, _}, _, St) ->
     #stream_acc{workers = Workers, ready = Ready, ring_opts = RingOpts} = St,
     case fabric_ring:node_down(NodeRef, Workers, Ready, RingOpts) of
         {ok, Workers1} ->
             {ok, St#stream_acc{workers = Workers1}};
         error ->
             {error, {nodedown, <<"progress not possible">>}}
     end;
 handle_stream_start({rexi_EXIT, Reason}, Worker, St) ->
     #stream_acc{
         workers = Workers,
         ready = Ready,
         replacements = Replacements,
         ring_opts = RingOpts
     } = St,
     case {fabric_ring:handle_error(Worker, Workers, Ready, RingOpts), Reason} of
         {{ok, Workers1}, _Reason} ->
             {ok, St#stream_acc{workers = Workers1}};
         {error, {maintenance_mode, _Node}} when Replacements /= undefined ->
             % Check if we have replacements for this range
             % and start the new workers if so.
             case lists:keytake(Worker#shard.range, 1, Replacements) of
                 {value, {_Range, WorkerReplacements}, NewReplacements} ->
                     FinalWorkers = lists:foldl(
                         fun(Repl, NewWorkers) ->
                             NewWorker = (St#stream_acc.start_fun)(Repl),
                             add_worker_to_watchdog(self(), NewWorker),
                             fabric_dict:store(NewWorker, waiting, NewWorkers)
                         end,
                         Workers,
                         WorkerReplacements
                     ),
                     % Assert that our replaced worker provides us
                     % the oppurtunity to make progress. Need to make sure
                     % to include already processed responses, since we are
                     % checking the full range and some workers have already
                     % responded and were removed from the workers list
                     ReadyWorkers = [{W, R} || {_, W, R} <- Ready],
                     AllWorkers = FinalWorkers ++ ReadyWorkers,
                     true = fabric_ring:is_progress_possible(AllWorkers),
                     NewRefs = fabric_dict:fetch_keys(FinalWorkers),
                     {new_refs, NewRefs, St#stream_acc{
                         workers = FinalWorkers,
                         replacements = NewReplacements
                     }};
                 false ->
                     % If we progress isn't possible and we don't have any
                     % replacements then we're dead in the water.
                     {error, {nodedown, <<"progress not possible">>}}
             end;
         {error, _} ->
             {error, fabric_util:error_info(Reason)}
     end;
 handle_stream_start(rexi_STREAM_INIT, {Worker, From}, St) ->
     #stream_acc{workers = Workers, ready = Ready, ring_opts = RingOpts} = St,
     case fabric_dict:lookup_element(Worker, Workers) of
         undefined ->
             % This worker lost the race with other partition copies, terminate
             rexi:stream_cancel(From),
             {ok, St};
         waiting ->
             case fabric_ring:handle_response(Worker, From, Workers, Ready, RingOpts) of
                 {ok, {Workers1, Ready1}} ->
                     % Don't have a full ring yet. Keep getting responses
                     {ok, St#stream_acc{workers = Workers1, ready = Ready1}};
                 {stop, Ready1} ->
                     % Have a full ring of workers. But don't ack the worker
                     % yet so they don't start sending us rows until we're ready
                     {stop, St#stream_acc{workers = [], ready = Ready1}}
             end
     end;
 handle_stream_start({ok, Error}, _, St) when Error == ddoc_updated; Error == insufficient_storage ->
     WaitingWorkers = [W || {W, _} <- St#stream_acc.workers],
     ReadyWorkers = [W || {W, _} <- St#stream_acc.ready],
     cleanup(WaitingWorkers ++ ReadyWorkers),
     {stop, Error};
 handle_stream_start(Else, _, _) ->
     exit({invalid_stream_start, Else}).

 % Spawn an auxiliary rexi worker watchdog which triggers cleanup if;
 % * nothing has been streamed for $timeout duration
 % * the coordinator (request) process is forceably killed and doesn't
 % get a chance to process its `after` fabric:clean/1 clause.

 % The watchdog is initially disabled. Clients can enable it by calling
 % enable_watchdog/0 before starting a fabric_stream. That client is responsible
 % for calling kick_watchdog/0 on activity to prevent the watchdog from acting.
 % If kick_watchdog/0 is not called at least once in each watchdog interval,
 % the stream coordinator is killed and the workers are cleaned up.
 spawn_worker_watchdog(Coordinator, Workers) ->
     case get(?WORKER_WATCHDOG) of
         undefined ->
             State0 = #watchdog_state{
                 coordinator = Coordinator,
                 workers = Workers
             },
             Enabled = get(?WATCHDOG_ENABLE),
             Pid = spawn(fun() ->
                 erlang:monitor(process, Coordinator),
                 State1 =
                     case Enabled of
                         true ->
                             reset_watchdog(State0);
                         undefined ->
                             State0
                     end,
                 watchdog_loop(State1)
             end),
             put(?WORKER_WATCHDOG, Pid),
             Pid;
         ExistingWatchdog ->
             ExistingWatchdog
     end.

 watchdog_loop(#watchdog_state{} = St) ->
     receive
         {?ADD_WORKER, Pid, Worker} when Pid == St#watchdog_state.coordinator ->
             Workers = St#watchdog_state.workers,
             watchdog_loop(St#watchdog_state{workers = [Worker | Workers]});
         ?WATCHDOG_KICK when ?WATCHDOG_IS_ENABLED ->
             watchdog_loop(St#watchdog_state{kicked = true});
         ?WATCHDOG_KICK ->
             watchdog_loop(St);
         ?WATCHDOG_TIMEOUT when ?WATCHDOG_IS_ENABLED, not ?WATCHDOG_WAS_KICKED ->
             exit(St#watchdog_state.coordinator, {shutdown, watchdog_fired}),
             watchdog_loop(St);
         ?WATCHDOG_TIMEOUT when ?WATCHDOG_IS_ENABLED ->
             watchdog_loop(reset_watchdog(St));
         ?WATCHDOG_TIMEOUT ->
             St;
         {'DOWN', _, _, Pid, _} when Pid == St#watchdog_state.coordinator ->
             fabric_util:cleanup(St#watchdog_state.workers)
     end.

 add_worker_to_watchdog(CoordinatorPid, Worker) ->
     send_to_watchdog({?ADD_WORKER, CoordinatorPid, Worker}).

 enable_watchdog() ->
     put(?WATCHDOG_ENABLE, true).

 reset_watchdog(#watchdog_state{} = St) ->
     case watchdog_timeout() of
         infinity ->
             St;
         Timeout ->
             TRef = erlang:send_after(Timeout, self(), ?WATCHDOG_TIMEOUT),
             St#watchdog_state{timer = TRef, kicked = false}
     end.

 kick_watchdog() ->
     case should_kick() of
         true ->
             send_to_watchdog(?WATCHDOG_KICK),
             update_last_kick();
         false ->
             ok
     end.


 should_kick() ->
     case watchdog_timeout() of
         infinity ->
             false;
         Timeout when is_integer(Timeout) ->
             case get(?WATCHDOG_LAST_KICK) of
                 undefined ->
                     true;
                 LastKick when is_integer(LastKick) ->
                     Now = erlang:monotonic_time(),
                     Delta = erlang:convert_time_unit(
                         Now - LastKick, native, millisecond
                     ),
                     Delta > Timeout
             end
     end.

 update_last_kick() ->
     put(?WATCHDOG_LAST_KICK, erlang:monotonic_time()).

 send_to_watchdog(Msg) ->
     case get(?WORKER_WATCHDOG) of
         WatchdogPid when is_pid(WatchdogPid) ->
             WatchdogPid ! Msg;
         _ ->
             ok
     end.

 watchdog_timeout() ->
     fabric_util:timeout("idle_stream", "60000").

 -ifdef(TEST).

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

 worker_watchdog_test_() ->
     {
         "Fabric spawn_worker_watchdog test",
         {
             setup,
             fun setup/0,
             fun teardown/1,
             fun(_) ->
                 [
                     should_clean_workers(),
                     does_not_fire_if_cleanup_called(),
                     should_clean_additional_worker_too()
                 ]
             end
         }
     }.

 should_clean_workers() ->
     ?_test(begin
         meck:reset(rexi),
         erase(?WORKER_WATCHDOG),
         Workers = [
             #shard{node = 'n1', ref = make_ref()},
             #shard{node = 'n2', ref = make_ref()}
         ],
         {Coord, _} = spawn_monitor(fun() ->
             receive
                 die -> ok
             end
         end),
         Watchdog = spawn_worker_watchdog(Coord, Workers),
         Ref = erlang:monitor(process, Watchdog),
         Coord ! die,
         receive
             {'DOWN', Ref, _, Watchdog, _} -> ok
         end,
         ?assertEqual(1, meck:num_calls(rexi, kill_all, 1))
     end).

 does_not_fire_if_cleanup_called() ->
     ?_test(begin
         meck:reset(rexi),
         erase(?WORKER_WATCHDOG),
         Workers = [
             #shard{node = 'n1', ref = make_ref()},
             #shard{node = 'n2', ref = make_ref()}
         ],
         {Coord, _} = spawn_monitor(fun() ->
             receive
                 die -> ok
             end
         end),
         Watchdog = spawn_worker_watchdog(Coord, Workers),
         Ref = erlang:monitor(process, Watchdog),
         cleanup(Workers),
         Coord ! die,
         receive
             {'DOWN', Ref, _, _, _} -> ok
         end,
         % 2 calls would be from cleanup/1 function. If cleanup process fired
         % too it would have been 4 calls total.
         ?assertEqual(1, meck:num_calls(rexi, kill_all, 1))
     end).

 should_clean_additional_worker_too() ->
     ?_test(begin
         meck:reset(rexi),
         erase(?WORKER_WATCHDOG),
         Workers = [
             #shard{node = 'n1', ref = make_ref()}
         ],
         {Coord, _} = spawn_monitor(fun() ->
             receive
                 die -> ok
             end
         end),
         Watchdog = spawn_worker_watchdog(Coord, Workers),
         add_worker_to_watchdog(Coord, #shard{node = 'n2', ref = make_ref()}),
         Ref = erlang:monitor(process, Watchdog),
         Coord ! die,
         receive
             {'DOWN', Ref, _, Watchdog, _} -> ok
         end,
         ?assertEqual(1, meck:num_calls(rexi, kill_all, 1))
     end).

 setup() ->
     ok = meck:expect(rexi, kill_all, fun(_) -> ok end).

 teardown(_) ->
     meck:unload().

 -endif.
	% 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(fabric_streams).

	-export([
	start/2,
	start/3,
	start/4,
	start/5,
	cleanup/1,
	enable_watchdog/0,
	kick_watchdog/0
	]).

	-include_lib("fabric/include/fabric.hrl").
	-include_lib("mem3/include/mem3.hrl").

	-define(WORKER_WATCHDOG, fabric_worker_watchdog).
	-define(ADD_WORKER, add_worker).
	-define(WATCHDOG_ENABLE, watchdog_enable).
	-define(WATCHDOG_TIMEOUT, watchdog_timeout).
	-define(WATCHDOG_KICK, watchdog_kick).
	-define(WATCHDOG_LAST_KICK, watchdog_last_kick).

	-define(WATCHDOG_WAS_KICKED, St#watchdog_state.kicked).
	-define(WATCHDOG_IS_ENABLED, St#watchdog_state.timer /= undefined).

	-record(watchdog_state, {
	coordinator,
	workers,
	timer,
	kicked = false
	}).

	start(Workers, Keypos) ->
	start(Workers, Keypos, undefined, undefined).

	start(Workers, Keypos, RingOpts) ->
	start(Workers, Keypos, undefined, undefined, RingOpts).

	start(Workers, Keypos, StartFun, Replacements) ->
	start(Workers, Keypos, StartFun, Replacements, []).

	start(Workers0, Keypos, StartFun, Replacements, RingOpts) ->
	Fun = fun handle_stream_start/3,
	Acc = #stream_acc{
	workers = fabric_dict:init(Workers0, waiting),
	ready = [],
	start_fun = StartFun,
	replacements = Replacements,
	ring_opts = RingOpts
	},
	spawn_worker_watchdog(self(), Workers0),
	Timeout = fabric_util:request_timeout(),
	case rexi_utils:recv(Workers0, Keypos, Fun, Acc, Timeout, infinity) of
	{ok, #stream_acc{ready = Workers}} ->
	AckedWorkers = fabric_dict:fold(
	fun(Worker, From, WorkerAcc) ->
	rexi:stream_start(From),
	[Worker \| WorkerAcc]
	end,
	[],
	Workers
	),
	{ok, AckedWorkers};
	Else ->
	Else
	end.

	cleanup(Workers) ->
	% Stop the auxiliary watchdog process as we got to the point where cleanup
	% happens in the regular fashion so we don't want to send 2x the number kill
	% messages
	case get(?WORKER_WATCHDOG) of
	WatchdogPid when is_pid(WatchdogPid) ->
	erase(?WORKER_WATCHDOG),
	exit(WatchdogPid, kill);
	_ ->
	ok
	end,
	fabric_util:cleanup(Workers).

	handle_stream_start({rexi_DOWN, _, {_, NodeRef}, _}, _, St) ->
	#stream_acc{workers = Workers, ready = Ready, ring_opts = RingOpts} = St,
	case fabric_ring:node_down(NodeRef, Workers, Ready, RingOpts) of
	{ok, Workers1} ->
	{ok, St#stream_acc{workers = Workers1}};
	error ->
	{error, {nodedown, <<"progress not possible">>}}
	end;
	handle_stream_start({rexi_EXIT, Reason}, Worker, St) ->
	#stream_acc{
	workers = Workers,
	ready = Ready,
	replacements = Replacements,
	ring_opts = RingOpts
	} = St,
	case {fabric_ring:handle_error(Worker, Workers, Ready, RingOpts), Reason} of
	{{ok, Workers1}, _Reason} ->
	{ok, St#stream_acc{workers = Workers1}};
	{error, {maintenance_mode, _Node}} when Replacements /= undefined ->
	% Check if we have replacements for this range
	% and start the new workers if so.
	case lists:keytake(Worker#shard.range, 1, Replacements) of
	{value, {_Range, WorkerReplacements}, NewReplacements} ->
	FinalWorkers = lists:foldl(
	fun(Repl, NewWorkers) ->
	NewWorker = (St#stream_acc.start_fun)(Repl),
	add_worker_to_watchdog(self(), NewWorker),
	fabric_dict:store(NewWorker, waiting, NewWorkers)
	end,
	Workers,
	WorkerReplacements
	),
	% Assert that our replaced worker provides us
	% the oppurtunity to make progress. Need to make sure
	% to include already processed responses, since we are
	% checking the full range and some workers have already
	% responded and were removed from the workers list
	ReadyWorkers = [{W, R} \|\| {_, W, R} <- Ready],
	AllWorkers = FinalWorkers ++ ReadyWorkers,
	true = fabric_ring:is_progress_possible(AllWorkers),
	NewRefs = fabric_dict:fetch_keys(FinalWorkers),
	{new_refs, NewRefs, St#stream_acc{
	workers = FinalWorkers,
	replacements = NewReplacements
	}};
	false ->
	% If we progress isn't possible and we don't have any
	% replacements then we're dead in the water.
	{error, {nodedown, <<"progress not possible">>}}
	end;
	{error, _} ->
	{error, fabric_util:error_info(Reason)}
	end;
	handle_stream_start(rexi_STREAM_INIT, {Worker, From}, St) ->
	#stream_acc{workers = Workers, ready = Ready, ring_opts = RingOpts} = St,
	case fabric_dict:lookup_element(Worker, Workers) of
	undefined ->
	% This worker lost the race with other partition copies, terminate
	rexi:stream_cancel(From),
	{ok, St};
	waiting ->
	case fabric_ring:handle_response(Worker, From, Workers, Ready, RingOpts) of
	{ok, {Workers1, Ready1}} ->
	% Don't have a full ring yet. Keep getting responses
	{ok, St#stream_acc{workers = Workers1, ready = Ready1}};
	{stop, Ready1} ->
	% Have a full ring of workers. But don't ack the worker
	% yet so they don't start sending us rows until we're ready
	{stop, St#stream_acc{workers = [], ready = Ready1}}
	end
	end;
	handle_stream_start({ok, Error}, _, St) when Error == ddoc_updated; Error == insufficient_storage ->
	WaitingWorkers = [W \|\| {W, _} <- St#stream_acc.workers],
	ReadyWorkers = [W \|\| {W, _} <- St#stream_acc.ready],
	cleanup(WaitingWorkers ++ ReadyWorkers),
	{stop, Error};
	handle_stream_start(Else, _, _) ->
	exit({invalid_stream_start, Else}).

	% Spawn an auxiliary rexi worker watchdog which triggers cleanup if;
	% * nothing has been streamed for $timeout duration
	% * the coordinator (request) process is forceably killed and doesn't
	% get a chance to process its `after` fabric:clean/1 clause.

	% The watchdog is initially disabled. Clients can enable it by calling
	% enable_watchdog/0 before starting a fabric_stream. That client is responsible
	% for calling kick_watchdog/0 on activity to prevent the watchdog from acting.
	% If kick_watchdog/0 is not called at least once in each watchdog interval,
	% the stream coordinator is killed and the workers are cleaned up.
	spawn_worker_watchdog(Coordinator, Workers) ->
	case get(?WORKER_WATCHDOG) of
	undefined ->
	State0 = #watchdog_state{
	coordinator = Coordinator,
	workers = Workers
	},
	Enabled = get(?WATCHDOG_ENABLE),
	Pid = spawn(fun() ->
	erlang:monitor(process, Coordinator),
	State1 =
	case Enabled of
	true ->
	reset_watchdog(State0);
	undefined ->
	State0
	end,
	watchdog_loop(State1)
	end),
	put(?WORKER_WATCHDOG, Pid),
	Pid;
	ExistingWatchdog ->
	ExistingWatchdog
	end.

	watchdog_loop(#watchdog_state{} = St) ->
	receive
	{?ADD_WORKER, Pid, Worker} when Pid == St#watchdog_state.coordinator ->
	Workers = St#watchdog_state.workers,
	watchdog_loop(St#watchdog_state{workers = [Worker \| Workers]});
	?WATCHDOG_KICK when ?WATCHDOG_IS_ENABLED ->
	watchdog_loop(St#watchdog_state{kicked = true});
	?WATCHDOG_KICK ->
	watchdog_loop(St);
	?WATCHDOG_TIMEOUT when ?WATCHDOG_IS_ENABLED, not ?WATCHDOG_WAS_KICKED ->
	exit(St#watchdog_state.coordinator, {shutdown, watchdog_fired}),
	watchdog_loop(St);
	?WATCHDOG_TIMEOUT when ?WATCHDOG_IS_ENABLED ->
	watchdog_loop(reset_watchdog(St));
	?WATCHDOG_TIMEOUT ->
	St;
	{'DOWN', _, _, Pid, _} when Pid == St#watchdog_state.coordinator ->
	fabric_util:cleanup(St#watchdog_state.workers)
	end.

	add_worker_to_watchdog(CoordinatorPid, Worker) ->
	send_to_watchdog({?ADD_WORKER, CoordinatorPid, Worker}).

	enable_watchdog() ->
	put(?WATCHDOG_ENABLE, true).

	reset_watchdog(#watchdog_state{} = St) ->
	case watchdog_timeout() of
	infinity ->
	St;
	Timeout ->
	TRef = erlang:send_after(Timeout, self(), ?WATCHDOG_TIMEOUT),
	St#watchdog_state{timer = TRef, kicked = false}
	end.

	kick_watchdog() ->
	case should_kick() of
	true ->
	send_to_watchdog(?WATCHDOG_KICK),
	update_last_kick();
	false ->
	ok
	end.


	should_kick() ->
	case watchdog_timeout() of
	infinity ->
	false;
	Timeout when is_integer(Timeout) ->
	case get(?WATCHDOG_LAST_KICK) of
	undefined ->
	true;
	LastKick when is_integer(LastKick) ->
	Now = erlang:monotonic_time(),
	Delta = erlang:convert_time_unit(
	Now - LastKick, native, millisecond
	),
	Delta > Timeout
	end
	end.

	update_last_kick() ->
	put(?WATCHDOG_LAST_KICK, erlang:monotonic_time()).

	send_to_watchdog(Msg) ->
	case get(?WORKER_WATCHDOG) of
	WatchdogPid when is_pid(WatchdogPid) ->
	WatchdogPid ! Msg;
	_ ->
	ok
	end.

	watchdog_timeout() ->
	fabric_util:timeout("idle_stream", "60000").

	-ifdef(TEST).

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

	worker_watchdog_test_() ->
	{
	"Fabric spawn_worker_watchdog test",
	{
	setup,
	fun setup/0,
	fun teardown/1,
	fun(_) ->
	[
	should_clean_workers(),
	does_not_fire_if_cleanup_called(),
	should_clean_additional_worker_too()
	]
	end
	}
	}.

	should_clean_workers() ->
	?_test(begin
	meck:reset(rexi),
	erase(?WORKER_WATCHDOG),
	Workers = [
	#shard{node = 'n1', ref = make_ref()},
	#shard{node = 'n2', ref = make_ref()}
	],
	{Coord, _} = spawn_monitor(fun() ->
	receive
	die -> ok
	end
	end),
	Watchdog = spawn_worker_watchdog(Coord, Workers),
	Ref = erlang:monitor(process, Watchdog),
	Coord ! die,
	receive
	{'DOWN', Ref, _, Watchdog, _} -> ok
	end,
	?assertEqual(1, meck:num_calls(rexi, kill_all, 1))
	end).

	does_not_fire_if_cleanup_called() ->
	?_test(begin
	meck:reset(rexi),
	erase(?WORKER_WATCHDOG),
	Workers = [
	#shard{node = 'n1', ref = make_ref()},
	#shard{node = 'n2', ref = make_ref()}
	],
	{Coord, _} = spawn_monitor(fun() ->
	receive
	die -> ok
	end
	end),
	Watchdog = spawn_worker_watchdog(Coord, Workers),
	Ref = erlang:monitor(process, Watchdog),
	cleanup(Workers),
	Coord ! die,
	receive
	{'DOWN', Ref, _, _, _} -> ok
	end,
	% 2 calls would be from cleanup/1 function. If cleanup process fired
	% too it would have been 4 calls total.
	?assertEqual(1, meck:num_calls(rexi, kill_all, 1))
	end).

	should_clean_additional_worker_too() ->
	?_test(begin
	meck:reset(rexi),
	erase(?WORKER_WATCHDOG),
	Workers = [
	#shard{node = 'n1', ref = make_ref()}
	],
	{Coord, _} = spawn_monitor(fun() ->
	receive
	die -> ok
	end
	end),
	Watchdog = spawn_worker_watchdog(Coord, Workers),
	add_worker_to_watchdog(Coord, #shard{node = 'n2', ref = make_ref()}),
	Ref = erlang:monitor(process, Watchdog),
	Coord ! die,
	receive
	{'DOWN', Ref, _, Watchdog, _} -> ok
	end,
	?assertEqual(1, meck:num_calls(rexi, kill_all, 1))
	end).

	setup() ->
	ok = meck:expect(rexi, kill_all, fun(_) -> ok end).

	teardown(_) ->
	meck:unload().

	-endif.