1%%%-------------------------------------------------------------------
2%%% File    : ibrowse_lb.erl
3%%% Author  : chandru <chandrashekhar.mullaparthi@t-mobile.co.uk>
4%%% Description :
5%%%
6%%% Created :  6 Mar 2008 by chandru <chandrashekhar.mullaparthi@t-mobile.co.uk>
7%%%-------------------------------------------------------------------
8-module(ibrowse_lb).
9-author(chandru).
10-behaviour(gen_server).
11%%--------------------------------------------------------------------
12%% Include files
13%%--------------------------------------------------------------------
14
15%%--------------------------------------------------------------------
16%% External exports
17-export([
18	 start_link/1,
19	 spawn_connection/6,
20         stop/1
21	]).
22
23%% gen_server callbacks
24-export([
25	 init/1,
26	 handle_call/3,
27	 handle_cast/2,
28	 handle_info/2,
29	 terminate/2,
30	 code_change/3
31	]).
32
33-record(state, {parent_pid,
34		ets_tid,
35		host,
36		port,
37		max_sessions,
38		max_pipeline_size,
39                proc_state
40               }).
41
42-include("ibrowse.hrl").
43
44%%====================================================================
45%% External functions
46%%====================================================================
47%%--------------------------------------------------------------------
48%% Function: start_link/0
49%% Description: Starts the server
50%%--------------------------------------------------------------------
51start_link(Args) ->
52    gen_server:start_link(?MODULE, Args, []).
53
54%%====================================================================
55%% Server functions
56%%====================================================================
57
58%%--------------------------------------------------------------------
59%% Function: init/1
60%% Description: Initiates the server
61%% Returns: {ok, State}          |
62%%          {ok, State, Timeout} |
63%%          ignore               |
64%%          {stop, Reason}
65%%--------------------------------------------------------------------
66init([Host, Port]) ->
67    process_flag(trap_exit, true),
68    Max_sessions = ibrowse:get_config_value({max_sessions, Host, Port}, 10),
69    Max_pipe_sz = ibrowse:get_config_value({max_pipeline_size, Host, Port}, 10),
70    put(my_trace_flag, ibrowse_lib:get_trace_status(Host, Port)),
71    put(ibrowse_trace_token, ["LB: ", Host, $:, integer_to_list(Port)]),
72    State = #state{parent_pid = whereis(ibrowse),
73		host = Host,
74		port = Port,
75		max_pipeline_size = Max_pipe_sz,
76	        max_sessions = Max_sessions},
77    State_1 = maybe_create_ets(State),
78    {ok, State_1}.
79
80spawn_connection(Lb_pid, Url,
81		 Max_sessions,
82		 Max_pipeline_size,
83		 SSL_options,
84		 Process_options)
85  when is_pid(Lb_pid),
86       is_record(Url, url),
87       is_integer(Max_pipeline_size),
88       is_integer(Max_sessions) ->
89    gen_server:call(Lb_pid,
90		    {spawn_connection, Url, Max_sessions, Max_pipeline_size, SSL_options, Process_options}).
91
92stop(Lb_pid) ->
93    case catch gen_server:call(Lb_pid, stop) of
94        {'EXIT', {timeout, _}} ->
95            exit(Lb_pid, kill);
96        ok ->
97            ok
98    end.
99%%--------------------------------------------------------------------
100%% Function: handle_call/3
101%% Description: Handling call messages
102%% Returns: {reply, Reply, State}          |
103%%          {reply, Reply, State, Timeout} |
104%%          {noreply, State}               |
105%%          {noreply, State, Timeout}      |
106%%          {stop, Reason, Reply, State}   | (terminate/2 is called)
107%%          {stop, Reason, State}            (terminate/2 is called)
108%%--------------------------------------------------------------------
109
110handle_call(stop, _From, #state{ets_tid = undefined} = State) ->
111    gen_server:reply(_From, ok),
112    {stop, normal, State};
113
114handle_call(stop, _From, #state{ets_tid = Tid} = State) ->
115    stop_all_conn_procs(Tid),
116    gen_server:reply(_From, ok),
117    {stop, normal, State};
118
119handle_call(_, _From, #state{proc_state = shutting_down} = State) ->
120    {reply, {error, shutting_down}, State};
121
122handle_call({spawn_connection, Url, Max_sess, Max_pipe, SSL_options, Process_options}, _From,
123	    State) ->
124    State_1   = maybe_create_ets(State),
125    Tid       = State_1#state.ets_tid,
126    Tid_size  = ets:info(Tid, size),
127    case Tid_size >= Max_sess of
128        true ->
129            Reply = find_best_connection(Tid, Max_pipe),
130            {reply, Reply, State_1#state{max_sessions      = Max_sess,
131                                         max_pipeline_size = Max_pipe}};
132        false ->
133            {ok, Pid} = ibrowse_http_client:start({Tid, Url, SSL_options}, Process_options),
134            Ts = os:timestamp(),
135            ets:insert(Tid, {{1, Ts, Pid}, []}),
136            {reply, {ok, {1, Ts, Pid}}, State_1#state{max_sessions      = Max_sess,
137						      max_pipeline_size = Max_pipe}}
138    end;
139
140handle_call(Request, _From, State) ->
141    Reply = {unknown_request, Request},
142    {reply, Reply, State}.
143
144%%--------------------------------------------------------------------
145%% Function: handle_cast/2
146%% Description: Handling cast messages
147%% Returns: {noreply, State}          |
148%%          {noreply, State, Timeout} |
149%%          {stop, Reason, State}            (terminate/2 is called)
150%%--------------------------------------------------------------------
151handle_cast(_Msg, State) ->
152    {noreply, State}.
153
154%%--------------------------------------------------------------------
155%% Function: handle_info/2
156%% Description: Handling all non call/cast messages
157%% Returns: {noreply, State}          |
158%%          {noreply, State, Timeout} |
159%%          {stop, Reason, State}            (terminate/2 is called)
160%%--------------------------------------------------------------------
161
162handle_info({trace, Bool}, #state{ets_tid = undefined} = State) ->
163    put(my_trace_flag, Bool),
164    {noreply, State};
165
166handle_info({trace, Bool}, #state{ets_tid = Tid} = State) ->
167    ets:foldl(fun({{_, Pid}, _}, Acc) when is_pid(Pid) ->
168		      catch Pid ! {trace, Bool},
169		      Acc;
170		 (_, Acc) ->
171		      Acc
172	      end, undefined, Tid),
173    put(my_trace_flag, Bool),
174    {noreply, State};
175
176handle_info(timeout, State) ->
177    %% We can't shutdown the process immediately because a request
178    %% might be in flight. So we first remove the entry from the
179    %% ibrowse_lb ets table, and then shutdown a couple of seconds
180    %% later
181    ets:delete(ibrowse_lb, {State#state.host, State#state.port}),
182    erlang:send_after(2000, self(), shutdown),
183    {noreply, State#state{proc_state = shutting_down}};
184
185handle_info(shutdown, State) ->
186    {stop, normal, State};
187
188handle_info(_Info, State) ->
189    {noreply, State}.
190
191%%--------------------------------------------------------------------
192%% Function: terminate/2
193%% Description: Shutdown the server
194%% Returns: any (ignored by gen_server)
195%%--------------------------------------------------------------------
196terminate(_Reason, #state{host = Host, port = Port, ets_tid = Tid} = _State) ->
197    catch ets:delete(ibrowse_lb, {Host, Port}),
198    stop_all_conn_procs(Tid),
199    ok.
200
201stop_all_conn_procs(Tid) ->
202    ets:foldl(fun({{_, _, Pid}, _}, Acc) ->
203                      ibrowse_http_client:stop(Pid),
204                      Acc
205              end, [], Tid).
206
207%%--------------------------------------------------------------------
208%% Func: code_change/3
209%% Purpose: Convert process state when code is changed
210%% Returns: {ok, NewState}
211%%--------------------------------------------------------------------
212code_change(_OldVsn, State, _Extra) ->
213    {ok, State}.
214
215%%--------------------------------------------------------------------
216%%% Internal functions
217%%--------------------------------------------------------------------
218find_best_connection(Tid, Max_pipe) ->
219    case ets:first(Tid) of
220        {Spec_size, Ts, Pid} = First when Spec_size < Max_pipe ->
221	    ets:delete(Tid, First),
222	    ets:insert(Tid, {{Spec_size + 1, Ts, Pid}, []}),
223	    {ok, First};
224        _ ->
225            {error, retry_later}
226    end.
227
228maybe_create_ets(#state{ets_tid = undefined, host = Host, port = Port} = State) ->
229    Tid = ets:new(ibrowse_lb, [public, ordered_set]),
230    ets:insert(ibrowse_lb, #lb_pid{host_port = {Host, Port}, pid = self(), ets_tid = Tid}),
231    State#state{ets_tid = Tid};
232maybe_create_ets(State) ->
233    State.
234