/edynamo/src/edynamo_handoff_receiver.erl

http://erlang-lab.googlecode.com/ · Erlang · 81 lines · 67 code · 10 blank · 4 comment · 4 complexity · 63f64c189ffad056ad405ddf9c8525ce MD5 · raw file

  1. %% copy riak-0.14.2 source to learn dynamo implementation and erlang
  2. %% author: lichuang
  3. %% copy source from riak_core_handoff_receiver.erl
  4. %% riak_handoff_receiver: incoming data handler for TCP-based handoff
  5. -module(edynamo_handoff_receiver).
  6. -include_lib("riak_core_handoff.hrl").
  7. -behaviour(gen_server2).
  8. -export([start_link/0,
  9. set_socket/2]).
  10. -export([init/1, handle_call/3, handle_cast/2, handle_info/2,
  11. terminate/2, code_change/3]).
  12. -record(state, {sock :: port(),
  13. partition :: non_neg_integer(),
  14. vnode_mod = edynamo_vnode:: module(),
  15. vnode :: pid(),
  16. count = 0 :: non_neg_integer()}).
  17. start_link() ->
  18. gen_server:start_link(?MODULE, [], []).
  19. set_socket(Pid, Socket) ->
  20. gen_server:call(Pid, {set_socket, Socket}).
  21. init([]) ->
  22. {ok, #state{}}.
  23. handle_call({set_socket, Socket}, _From, State) ->
  24. inet:setopts(Socket, [{active, once}, {packet, 4}, {header, 1}]),
  25. {reply, ok, State#state { sock = Socket }}.
  26. handle_info({tcp_closed, _Socket}, State=#state{partition = Partition, count = Count}) ->
  27. error_logger:info_msg("Handoff receiver for partition ~p exiting after processing ~p"
  28. " objects~n", [Partition, Count]),
  29. {stop, normal, State};
  30. handle_info({tcp_error, _Socket, _Reason}, State=#state{partition=Partition,count=Count}) ->
  31. error_logger:info_msg("Handoff receiver for partition ~p exiting after processing ~p"
  32. " objects~n", [Partition, Count]),
  33. {stop, normal, State};
  34. handle_info({tcp, Socket, Data}, State) ->
  35. [MsgType|MsgData] = Data,
  36. case catch(process_message(MsgType, MsgData, State)) of
  37. {'EXIT', Reason} ->
  38. error_logger:error_msg("Handoff receiver for partition ~p exiting abnormally after "
  39. "processing ~p objects: ~p\n", [State#state.partition, State#state.count, Reason]),
  40. {stop, normal, State};
  41. NewState when is_record(NewState, state) ->
  42. inet:setopts(Socket, [{active, once}]),
  43. {noreply, NewState}
  44. end.
  45. process_message(?PT_MSG_INIT, MsgData, State=#state{vnode_mod=VNodeMod}) ->
  46. <<Partition:160/integer>> = MsgData,
  47. error_logger:info_msg("Receiving handoff data for partition ~p:~p~n", [VNodeMod, Partition]),
  48. {ok, VNode} = edynamo_vnode_master:get_vnode_pid(Partition, VNodeMod),
  49. State#state{partition=Partition, vnode=VNode};
  50. process_message(?PT_MSG_OBJ, MsgData, State=#state{vnode=VNode, count = Count}) ->
  51. Msg = {handoff_data, MsgData},
  52. gen_fsm:sync_send_all_state_event(VNode, Msg, 60000),
  53. State#state{count = Count + 1};
  54. process_message(?PT_MSG_OLDSYNC, MsgData, State=#state{sock = Socket}) ->
  55. gen_tcp:send(Socket, <<?PT_MSG_OLDSYNC:8, "sync">>),
  56. <<VNodeModBin/binary>> = MsgData,
  57. VNodeMod = binary_to_atom(VNodeModBin, utf8),
  58. State#state{vnode_mod = VNodeMod};
  59. process_message(?PT_MSG_SYNC, _MsgData, State=#state{sock=Socket}) ->
  60. gen_tcp:send(Socket, <<?PT_MSG_SYNC:8, "sync">>),
  61. State;
  62. process_message(?PT_MSG_CONFIGURE, MsgData, State) ->
  63. ConfProps = binary_to_term(MsgData),
  64. State#state{vnode_mod=proplists:get_value(vnode_mod, ConfProps),
  65. partition=proplists:get_value(partition, ConfProps)};
  66. process_message(_, _MsgData, State=#state{sock=Socket}) ->
  67. gen_tcp:send(Socket, <<255:8,"unknown_msg">>),
  68. State.
  69. handle_cast(_Msg, State) -> {noreply, State}.
  70. terminate(_Reason, _State) -> ok.
  71. code_change(_OldVsn, State, _Extra) -> {ok, State}.