%% ------------------------------------------------------------------- %% %% Copyright (c) 2013 Basho Technologies, Inc. All Rights Reserved. %% %% This file is provided to you 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. %% %% ------------------------------------------------------------------- %% @doc %% This module implements the sidejob_worker logic used by all worker %% processes created to manage a sidejob resource. This code emulates %% the gen_server API, wrapping a provided user-specified module which %% implements the gen_server behavior. %% %% The primary purpose of this module is updating the usage information %% published in a given resource's ETS table, such that capacity limiting %% operates correctly. The sidejob_worker also cooperates with a given %% {@link sidejob_resource_stats} server to maintain statistics about a %% given resource. %% %% By default, a sidejob_worker calculates resource usage based on message %% queue size. However, the user-specified module can also choose to %% implement the `current_usage/1' and `rate/1' callbacks to change how %% usage is calculated. An example is the {@link sidejob_supervisor} module %% which reports usage as: queue size + num_children. -module(sidejob_worker). -behaviour(gen_server). %% API -export([start_link/6, reg_name/1, reg_name/2, workers/2]). %% gen_server callbacks -export([init/1, handle_call/3, handle_cast/2, handle_info/2, terminate/2, code_change/3]). -record(state, {id :: non_neg_integer(), ets :: term(), width :: pos_integer(), limit :: pos_integer(), reporter :: term(), mod :: module(), modstate :: term(), usage :: custom | default, last_mq_len = 0 :: non_neg_integer(), enqueue = 0 :: non_neg_integer(), dequeue = 0 :: non_neg_integer()}). %%%=================================================================== %%% API %%%=================================================================== reg_name(Id) -> IdBin = list_to_binary(integer_to_list(Id)), binary_to_atom(<<"sidejob_worker_", IdBin/binary>>, latin1). reg_name(Name, Id) when is_atom(Name) -> reg_name(atom_to_binary(Name, latin1), Id); reg_name(NameBin, Id) -> WorkerName = iolist_to_binary([NameBin, "_", integer_to_list(Id)]), binary_to_atom(WorkerName, latin1). workers(Name, Count) -> NameBin = atom_to_binary(Name, latin1), [reg_name(NameBin, Id) || Id <- lists:seq(1,Count)]. start_link(RegName, ResName, Id, ETS, StatsName, Mod) -> gen_server:start_link({local, RegName}, ?MODULE, [ResName, Id, ETS, StatsName, Mod], []). %%%=================================================================== %%% gen_server callbacks %%%=================================================================== init([ResName, Id, ETS, StatsName, Mod]) -> %% TODO: Add ability to pass args case Mod:init([ResName]) of {ok, ModState} -> Exports = proplists:get_value(exports, Mod:module_info()), Usage = case lists:member({current_usage, 1}, Exports) of true -> custom; false -> default end, schedule_tick(), Width = ResName:width(), Limit = ResName:limit(), State = #state{id=Id, ets=ETS, mod=Mod, modstate=ModState, usage=Usage, width=Width, limit=Limit, reporter=StatsName}, ets:insert(ETS, [{usage,0}, {full,0}]), {ok, State}; Other -> Other end. handle_call(Request, From, State=#state{mod=Mod, modstate=ModState}) -> Result = Mod:handle_call(Request, From, ModState), {Pos, ModState2} = case Result of {reply,_Reply,NewState} -> {3, NewState}; {reply,_Reply,NewState,hibernate} -> {3, NewState}; {reply,_Reply,NewState,_Timeout} -> {3, NewState}; {noreply,NewState} -> {2, NewState}; {noreply,NewState,hibernate} -> {2, NewState}; {noreply,NewState,_Timeout} -> {2, NewState}; {stop,_Reason,_Reply,NewState} -> {4, NewState}; {stop,_Reason,NewState} -> {3, NewState} end, State2 = State#state{modstate=ModState2}, State3 = update_rate(update_usage(State2)), Return = setelement(Pos, Result, State3), Return. handle_cast(Request, State=#state{mod=Mod, modstate=ModState}) -> Result = Mod:handle_cast(Request, ModState), {Pos, ModState2} = case Result of {noreply,NewState} -> {2, NewState}; {noreply,NewState,hibernate} -> {2, NewState}; {noreply,NewState,_Timeout} -> {2, NewState}; {stop,_Reason,NewState} -> {3, NewState} end, State2 = State#state{modstate=ModState2}, State3 = update_rate(update_usage(State2)), Return = setelement(Pos, Result, State3), Return. handle_info('$sidejob_worker_tick', State) -> State2 = tick(State), schedule_tick(), {noreply, State2}; handle_info(Info, State=#state{mod=Mod, modstate=ModState}) -> Result = Mod:handle_info(Info, ModState), {Pos, ModState2} = case Result of {noreply,NewState} -> {2, NewState}; {noreply,NewState,hibernate} -> {2, NewState}; {noreply,NewState,_Timeout} -> {2, NewState}; {stop,_Reason,NewState} -> {3, NewState} end, State2 = State#state{modstate=ModState2}, State3 = update_rate(update_usage(State2)), Return = setelement(Pos, Result, State3), Return. terminate(_Reason, _State) -> ok. code_change(_OldVsn, State, _Extra) -> {ok, State}. %%%=================================================================== %%% Internal functions %%%=================================================================== schedule_tick() -> erlang:send_after(1000, self(), '$sidejob_worker_tick'). tick(State=#state{id=Id, reporter=Reporter}) -> Usage = current_usage(State), {In, Out, State2} = current_rate(State), sidejob_resource_stats:report(Reporter, Id, Usage, In, Out), State2. update_usage(State=#state{ets=ETS, width=Width, limit=Limit}) -> Usage = current_usage(State), Full = case Usage >= (Limit div Width) of true -> 1; false -> 0 end, ets:insert(ETS, [{usage, Usage}, {full, Full}]), State. current_usage(#state{usage=default}) -> {message_queue_len, Len} = process_info(self(), message_queue_len), Len; current_usage(#state{usage=custom, mod=Mod, modstate=ModState}) -> Mod:current_usage(ModState). update_rate(State=#state{usage=custom}) -> %% Assume this is updated internally in the custom module State; update_rate(State=#state{usage=default, last_mq_len=LastLen}) -> {message_queue_len, Len} = process_info(self(), message_queue_len), Enqueue = Len - LastLen + 1, Dequeue = State#state.dequeue + 1, State#state{enqueue=Enqueue, dequeue=Dequeue}. %% TODO: Probably should rename since it resets rate current_rate(State=#state{usage=default, enqueue=Enqueue, dequeue=Dequeue}) -> State2 = State#state{enqueue=0, dequeue=0}, {Enqueue, Dequeue, State2}; current_rate(State=#state{usage=custom, mod=Mod, modstate=ModState}) -> {Enqueue, Dequeue, ModState2} = Mod:rate(ModState), State2 = State#state{modstate=ModState2}, {Enqueue, Dequeue, State2}.