src/ioq_osq.erl - couchdb-ioq - 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(ioq_osq).
 -behaviour(gen_server).
 -vsn(1).
 -export([init/1, handle_call/3, handle_cast/2, handle_info/2, terminate/2,
     code_change/3]).

 -export([start_link/0, call/3]).

 -record(channel, {
     name,
     q = queue:new()
 }).

 -record(state, {
     reqs = [],
     min = 2,
     max = 6,
     global_max = 15,
     channels = queue:new()
 }).

 %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% WARNING %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
 %% This server relies on the internal structure of the channels queue as a   %%
 %% {list(), list()} to do in-place modifications of some elements.  We are   %%
 %% "running on thin ice", as it were.                                        %%
 %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

 start_link() ->
     gen_server:start_link({local, ?MODULE}, ?MODULE, [], []).


 call(Pid, Msg, Priority) ->
     Now = erlang:monotonic_time(),
     Reply = gen_server:call(ioq_osq, {rlimit, nil, Priority, Now}, infinity),
     try
         gen_server:call(Pid, Msg, infinity)
     after
         whereis(ioq_osq) ! Reply
     end.


 init([]) ->
     ets:new(osq_counters, [named_table]),
     St = #state{},
     {ok, St#state{
         min = threshold("minimum", St#state.min),
         max = threshold("maximum", St#state.max),
         global_max = threshold("global_maximum", St#state.global_max)
     }}.

 handle_call({set_minimum, C}, _From, State) when is_integer(C), C > 0 ->
     {reply, ok, State#state{min = C}};

 handle_call({set_maximum, C}, _From, State) when is_integer(C), C > 0 ->
     {reply, ok, State#state{max = C}};

 handle_call({set_global_maximum, C}, _From, State) when is_integer(C), C > 0 ->
     {reply, ok, State#state{global_max = C}};

 handle_call(get_queue_depths, _From, State) ->
     Channels = [{N, queue:len(Q)} || #channel{name=N, q=Q}
         <- queue:to_list(State#state.channels)],
     {reply, Channels, State};

 handle_call(get_requests, _From, State) ->
     {reply, State#state.reqs, State};

 handle_call({_, _, {interactive, Shard}, _} = Req, From, State) ->
     {noreply, enqueue_channel(channel_name(Shard), {Req, From}, State)};

 handle_call({_, _, {view_update, Shard, _}, _} = Req, From, State) ->
     {noreply, enqueue_channel(channel_name(Shard), {Req, From}, State)};

 handle_call({Fd, _, _, _} = Req, From, State) when is_pid(Fd) ->
     {noreply, enqueue_channel(other, {Req, From}, State)};

 handle_call({rlimit, _, _, _} = Req, From, State) ->
     {noreply, enqueue_channel(other, {Req, From}, State)};

 handle_call(_Msg, _From, State) ->
     {reply, ignored, State}.

 handle_cast(_Msg, State) ->
     {noreply, State}.

 handle_info({Ref, Reply}, #state{reqs = Reqs} = State) ->
     case lists:keyfind(Ref, 3, Reqs) of
     {_, notify, Ref} ->
         erlang:demonitor(Ref, [flush]),
         Reqs2 = lists:keydelete(Ref, 3, Reqs),
         {noreply, make_next_request(State#state{reqs = Reqs2})};
     {_, From, Ref} ->
         erlang:demonitor(Ref, [flush]),
         gen_server:reply(From, Reply),
         Reqs2 = lists:keydelete(Ref, 3, Reqs),
         {noreply, make_next_request(State#state{reqs = Reqs2})};
     false ->
         {noreply, State}
     end;

 handle_info({'DOWN', Ref, _, _, Reason}, #state{reqs = Reqs} = State) ->
     case lists:keyfind(Ref, 3, Reqs) of
     {_, notify, Ref} ->
         Reqs2 = lists:keydelete(Ref, 3, Reqs),
         {noreply, make_next_request(State#state{reqs = Reqs2})};
     {_, From, Ref} ->
         gen_server:reply(From, {'EXIT', Reason}),
         Reqs2 = lists:keydelete(Ref, 3, Reqs),
         {noreply, make_next_request(State#state{reqs = Reqs2})};
     false ->
         {noreply, State}
     end;

 handle_info(_Info, State) ->
     {noreply, State}.

 terminate(_Reason, _State) ->
     ok.

 code_change(_OldVsn, State, _Extra) ->
     {ok, State}.

 channel_name(Shard) ->
     try re:split(Shard, "/") of
     [<<"shards">>, _, <<"heroku">>, AppId | _] ->
         <<AppId/binary, ".heroku">>;
     [<<"shards">>, _, Account | _] ->
         Account;
     _ ->
         other
     catch _:_ ->
         other
     end.

 find_channel(Account, {A, B}) ->
     case lists:keyfind(Account, 2, A) of
     false ->
         case lists:keyfind(Account, 2, B) of
         false ->
             {new, #channel{name = Account}};
         #channel{} = Channel ->
             {2, Channel}
         end;
     #channel{} = Channel ->
         {1, Channel}
     end.

 update_channel(#channel{q = Q} = Ch, Req) ->
     Ch#channel{q = queue:in(Req, Q)}.

 enqueue_channel(Account, Req, #state{channels = Q} = State) ->
     NewState = case find_channel(Account, Q) of
     {new, Channel0} ->
         State#state{channels = queue:in(update_channel(Channel0, Req), Q)};
     {Elem, Channel0} ->
         Channel = update_channel(Channel0, Req),
         % the channel already exists in the queue - update it in place
         L = element(Elem, Q),
         NewQ = setelement(Elem, Q, lists:keyreplace(Account, 2, L, Channel)),
         State#state{channels = NewQ}
     end,
     maybe_submit_request(NewState).

 maybe_submit_request(#state{global_max=C, reqs=R} = St) when length(R) < C ->
     make_next_request(St);
 maybe_submit_request(#state{min = Min} = State) ->
     % look for a channel which hasn't reached the minimum yet
     make_next_request(State, Min).

 make_next_request(#state{max = Max} = State) ->
     % default behavior, look for a channel not yet maxed out
     make_next_request(State, Max).

 make_next_request(#state{channels = Channels, reqs = R} = State, Threshold) ->
     case next_unblocked_channel(Channels, R, Threshold, queue:new()) of
     {#channel{name = Name, q = Q} = Ch, OutChannels} ->
         {{value, Item}, NewQ} = queue:out(Q),
         case queue:is_empty(NewQ) of
         true ->
             NewCh = OutChannels;
         false ->
             NewCh = queue:in(Ch#channel{q = NewQ}, OutChannels)
         end,
         submit_request(Name, Item, State#state{channels = NewCh});
     {nil, OutQ} ->
         % everybody is using their allotted slots, try again later
         State#state{channels = OutQ}
     end.

 next_unblocked_channel(InQ, Reqs, Max, OutQ)  ->
     case queue:out(InQ) of
     {empty, _} -> % all channels blocked
         {nil, OutQ};
     {{value, #channel{name=Name} = Channel}, NewQ} ->
         case length([1 || {N, _, _} <- Reqs, N =:= Name]) >= Max of
         true -> % channel is blocked, keep searching
             next_unblocked_channel(NewQ, Reqs, Max, queue:in(Channel, OutQ));
         false ->
             {Channel, queue:join(NewQ, OutQ)}
         end
     end.


 submit_request(Channel, {{rlimit,_,Pri,T0}, From}, #state{reqs=Reqs} = State) ->
     % rlimit fd means that we'll get a response back
     % from the pid after it performs the call on its
     % own
     Ref = erlang:monitor(process, element(1, From)),
     gen_server:reply(From, {Ref, nil}),
     record_stats(Channel, Pri, T0),
     State#state{reqs = [{Channel, notify, Ref} | Reqs]};

 submit_request(Channel, {{Fd,Call,Pri,T0}, From}, #state{reqs=Reqs} = State) ->
     % make the request
     Ref = erlang:monitor(process, Fd),
     Fd ! {'$gen_call', {self(), Ref}, Call},
     record_stats(Channel, Pri, T0),
     State#state{reqs = [{Channel, From, Ref} | Reqs]}.

 record_stats(Channel, Pri, T0) ->
     IOClass = if is_tuple(Pri) -> element(1, Pri); true -> Pri end,
     Now = erlang:monotonic_time(),
     Latency = erlang:convert_time_unit(
         Now - T0, native, millisecond),
     catch couch_stats:increment_counter([couchdb, io_queue, IOClass]),
     catch couch_stats:increment_counter([couchdb, io_queue, osproc]),
     catch couch_stats:update_histogram([couchdb, io_queue, latency], Latency),
     update_counter(Channel, IOClass, osproc).

 update_counter(Channel, IOClass, RW) ->
     try ets:update_counter(osq_counters, {Channel, IOClass, RW}, 1)
     catch error:badarg ->
         ets:insert(osq_counters, {{Channel, IOClass, RW}, 1})
     end.

 threshold(Name, Default) ->
     try list_to_integer(config:get("osq", Name)) of
     C when C > 0->
         C;
     _ ->
         Default
     catch _:_ ->
         Default
     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(ioq_osq).
	-behaviour(gen_server).
	-vsn(1).
	-export([init/1, handle_call/3, handle_cast/2, handle_info/2, terminate/2,
	code_change/3]).

	-export([start_link/0, call/3]).

	-record(channel, {
	name,
	q = queue:new()
	}).

	-record(state, {
	reqs = [],
	min = 2,
	max = 6,
	global_max = 15,
	channels = queue:new()
	}).

	%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% WARNING %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
	%% This server relies on the internal structure of the channels queue as a %%
	%% {list(), list()} to do in-place modifications of some elements. We are %%
	%% "running on thin ice", as it were. %%
	%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

	start_link() ->
	gen_server:start_link({local, ?MODULE}, ?MODULE, [], []).


	call(Pid, Msg, Priority) ->
	Now = erlang:monotonic_time(),
	Reply = gen_server:call(ioq_osq, {rlimit, nil, Priority, Now}, infinity),
	try
	gen_server:call(Pid, Msg, infinity)
	after
	whereis(ioq_osq) ! Reply
	end.


	init([]) ->
	ets:new(osq_counters, [named_table]),
	St = #state{},
	{ok, St#state{
	min = threshold("minimum", St#state.min),
	max = threshold("maximum", St#state.max),
	global_max = threshold("global_maximum", St#state.global_max)
	}}.

	handle_call({set_minimum, C}, _From, State) when is_integer(C), C > 0 ->
	{reply, ok, State#state{min = C}};

	handle_call({set_maximum, C}, _From, State) when is_integer(C), C > 0 ->
	{reply, ok, State#state{max = C}};

	handle_call({set_global_maximum, C}, _From, State) when is_integer(C), C > 0 ->
	{reply, ok, State#state{global_max = C}};

	handle_call(get_queue_depths, _From, State) ->
	Channels = [{N, queue:len(Q)} \|\| #channel{name=N, q=Q}
	<- queue:to_list(State#state.channels)],
	{reply, Channels, State};

	handle_call(get_requests, _From, State) ->
	{reply, State#state.reqs, State};

	handle_call({_, _, {interactive, Shard}, _} = Req, From, State) ->
	{noreply, enqueue_channel(channel_name(Shard), {Req, From}, State)};

	handle_call({_, _, {view_update, Shard, _}, _} = Req, From, State) ->
	{noreply, enqueue_channel(channel_name(Shard), {Req, From}, State)};

	handle_call({Fd, _, _, _} = Req, From, State) when is_pid(Fd) ->
	{noreply, enqueue_channel(other, {Req, From}, State)};

	handle_call({rlimit, _, _, _} = Req, From, State) ->
	{noreply, enqueue_channel(other, {Req, From}, State)};

	handle_call(_Msg, _From, State) ->
	{reply, ignored, State}.

	handle_cast(_Msg, State) ->
	{noreply, State}.

	handle_info({Ref, Reply}, #state{reqs = Reqs} = State) ->
	case lists:keyfind(Ref, 3, Reqs) of
	{_, notify, Ref} ->
	erlang:demonitor(Ref, [flush]),
	Reqs2 = lists:keydelete(Ref, 3, Reqs),
	{noreply, make_next_request(State#state{reqs = Reqs2})};
	{_, From, Ref} ->
	erlang:demonitor(Ref, [flush]),
	gen_server:reply(From, Reply),
	Reqs2 = lists:keydelete(Ref, 3, Reqs),
	{noreply, make_next_request(State#state{reqs = Reqs2})};
	false ->
	{noreply, State}
	end;

	handle_info({'DOWN', Ref, _, _, Reason}, #state{reqs = Reqs} = State) ->
	case lists:keyfind(Ref, 3, Reqs) of
	{_, notify, Ref} ->
	Reqs2 = lists:keydelete(Ref, 3, Reqs),
	{noreply, make_next_request(State#state{reqs = Reqs2})};
	{_, From, Ref} ->
	gen_server:reply(From, {'EXIT', Reason}),
	Reqs2 = lists:keydelete(Ref, 3, Reqs),
	{noreply, make_next_request(State#state{reqs = Reqs2})};
	false ->
	{noreply, State}
	end;

	handle_info(_Info, State) ->
	{noreply, State}.

	terminate(_Reason, _State) ->
	ok.

	code_change(_OldVsn, State, _Extra) ->
	{ok, State}.

	channel_name(Shard) ->
	try re:split(Shard, "/") of
	[<<"shards">>, _, <<"heroku">>, AppId \| _] ->
	<<AppId/binary, ".heroku">>;
	[<<"shards">>, _, Account \| _] ->
	Account;
	_ ->
	other
	catch _:_ ->
	other
	end.

	find_channel(Account, {A, B}) ->
	case lists:keyfind(Account, 2, A) of
	false ->
	case lists:keyfind(Account, 2, B) of
	false ->
	{new, #channel{name = Account}};
	#channel{} = Channel ->
	{2, Channel}
	end;
	#channel{} = Channel ->
	{1, Channel}
	end.

	update_channel(#channel{q = Q} = Ch, Req) ->
	Ch#channel{q = queue:in(Req, Q)}.

	enqueue_channel(Account, Req, #state{channels = Q} = State) ->
	NewState = case find_channel(Account, Q) of
	{new, Channel0} ->
	State#state{channels = queue:in(update_channel(Channel0, Req), Q)};
	{Elem, Channel0} ->
	Channel = update_channel(Channel0, Req),
	% the channel already exists in the queue - update it in place
	L = element(Elem, Q),
	NewQ = setelement(Elem, Q, lists:keyreplace(Account, 2, L, Channel)),
	State#state{channels = NewQ}
	end,
	maybe_submit_request(NewState).

	maybe_submit_request(#state{global_max=C, reqs=R} = St) when length(R) < C ->
	make_next_request(St);
	maybe_submit_request(#state{min = Min} = State) ->
	% look for a channel which hasn't reached the minimum yet
	make_next_request(State, Min).

	make_next_request(#state{max = Max} = State) ->
	% default behavior, look for a channel not yet maxed out
	make_next_request(State, Max).

	make_next_request(#state{channels = Channels, reqs = R} = State, Threshold) ->
	case next_unblocked_channel(Channels, R, Threshold, queue:new()) of
	{#channel{name = Name, q = Q} = Ch, OutChannels} ->
	{{value, Item}, NewQ} = queue:out(Q),
	case queue:is_empty(NewQ) of
	true ->
	NewCh = OutChannels;
	false ->
	NewCh = queue:in(Ch#channel{q = NewQ}, OutChannels)
	end,
	submit_request(Name, Item, State#state{channels = NewCh});
	{nil, OutQ} ->
	% everybody is using their allotted slots, try again later
	State#state{channels = OutQ}
	end.

	next_unblocked_channel(InQ, Reqs, Max, OutQ) ->
	case queue:out(InQ) of
	{empty, _} -> % all channels blocked
	{nil, OutQ};
	{{value, #channel{name=Name} = Channel}, NewQ} ->
	case length([1 \|\| {N, _, _} <- Reqs, N =:= Name]) >= Max of
	true -> % channel is blocked, keep searching
	next_unblocked_channel(NewQ, Reqs, Max, queue:in(Channel, OutQ));
	false ->
	{Channel, queue:join(NewQ, OutQ)}
	end
	end.


	submit_request(Channel, {{rlimit,_,Pri,T0}, From}, #state{reqs=Reqs} = State) ->
	% rlimit fd means that we'll get a response back
	% from the pid after it performs the call on its
	% own
	Ref = erlang:monitor(process, element(1, From)),
	gen_server:reply(From, {Ref, nil}),
	record_stats(Channel, Pri, T0),
	State#state{reqs = [{Channel, notify, Ref} \| Reqs]};

	submit_request(Channel, {{Fd,Call,Pri,T0}, From}, #state{reqs=Reqs} = State) ->
	% make the request
	Ref = erlang:monitor(process, Fd),
	Fd ! {'$gen_call', {self(), Ref}, Call},
	record_stats(Channel, Pri, T0),
	State#state{reqs = [{Channel, From, Ref} \| Reqs]}.

	record_stats(Channel, Pri, T0) ->
	IOClass = if is_tuple(Pri) -> element(1, Pri); true -> Pri end,
	Now = erlang:monotonic_time(),
	Latency = erlang:convert_time_unit(
	Now - T0, native, millisecond),
	catch couch_stats:increment_counter([couchdb, io_queue, IOClass]),
	catch couch_stats:increment_counter([couchdb, io_queue, osproc]),
	catch couch_stats:update_histogram([couchdb, io_queue, latency], Latency),
	update_counter(Channel, IOClass, osproc).

	update_counter(Channel, IOClass, RW) ->
	try ets:update_counter(osq_counters, {Channel, IOClass, RW}, 1)
	catch error:badarg ->
	ets:insert(osq_counters, {{Channel, IOClass, RW}, 1})
	end.

	threshold(Name, Default) ->
	try list_to_integer(config:get("osq", Name)) of
	C when C > 0->
	C;
	_ ->
	Default
	catch _:_ ->
	Default
	end.