/
riak_mongo_server.erl
96 lines (77 loc) · 2.74 KB
/
riak_mongo_server.erl
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
%%
%% This file is part of riak_mongo
%%
%% Copyright (c) 2012 by Pavlo Baron (pb at pbit dot org)
%%
%% 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.
%%
%% @author Pavlo Baron <pbat pbit dot org>
%% @doc This is the TCP server of riak_mongo
%% @copyright 2012 Pavlo Baron
-module(riak_mongo_server).
-export([start_link/0]).
-export([init/2, handle_call/3, handle_cast/2, handle_info/2]).
-export([terminate/2, sock_opts/0, new_connection/4]).
-export([add_listener/1]).
-behavior(gen_nb_server).
start_link() ->
gen_nb_server:start_link(?MODULE, []).
add_listener([IpAddr, Port]) ->
gen_server:call(whereis(?MODULE), {add_listener, IpAddr, Port}).
remove_listener(Pid, IpAddr, Port) ->
gen_server:call(Pid, {remove_listener, IpAddr, Port}).
init([], State) ->
register(?MODULE, self()),
{ok, gen_nb_server:store_cb_state([], State)}.
handle_call({add_listener, IpAddr, Port}, _From, State) ->
[] = gen_nb_server:get_cb_state(State),
case gen_nb_server:add_listen_socket({IpAddr, Port}, State) of
{ok, State1} ->
{reply, ok, State1};
Error ->
{reply, Error, State}
end;
handle_call({remove_listener, IpAddr, Port}, _From, State) ->
case gen_nb_server:remove_listen_socket({IpAddr, Port}, State) of
{ok, State1} ->
{reply, ok, State1};
Error ->
{reply, Error, State}
end;
handle_call(_Msg, _From, State) ->
{reply, ignored, State}.
handle_cast(_Msg, State) ->
{noreply, State}.
handle_info({tcp, Sock, Data}, State) ->
Me = self(),
P = spawn(fun() -> worker(Me, Sock, Data) end),
gen_tcp:controlling_process(Sock, P),
{noreply, State};
handle_info(_Msg, State) ->
{noreply, State}.
terminate(_Reason, _State) ->
ok.
sock_opts() ->
[binary, {active, once}, {packet, 0}].
new_connection(_IpAddr, _Port, Sock, State) ->
Me = self(),
P = spawn(fun() -> worker(Me, Sock) end),
gen_tcp:controlling_process(Sock, P),
{ok, State}.
worker(Owner, Sock) ->
inet:setopts(Sock, [{active, once}]),
gen_tcp:controlling_process(Sock, Owner).
worker(Owner, Sock, Data) ->
gen_tcp:send(Sock, wire_protocol:process_data(Sock, Data)),
inet:setopts(Sock, [{active, once}]),
gen_tcp:controlling_process(Sock, Owner).