%% @author Northscale <info@northscale.com>
%% @copyright 2009 NorthScale, Inc.
%%
%% 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(ns_pubsub).

-include("ns_common.hrl").

-behaviour(gen_event).

-include_lib("eunit/include/eunit.hrl").

-record(state, {func, func_state}).

%% API
-export([subscribe_link/1, subscribe_link/2, subscribe_link/3, unsubscribe/1]).

%% gen_event callbacks
-export([code_change/3, init/1, handle_call/2, handle_event/2, handle_info/2,
         terminate/2]).

%% called by proc_lib:start from subscribe_link/3
-export([do_subscribe_link/4]).


%%
%% API
%%
subscribe_link(Name) ->
    subscribe_link(Name, msg_fun(self()), ignored).

subscribe_link(Name, Fun) ->
    subscribe_link(
      Name,
      fun (Event, State) ->
              Fun(Event),
              State
      end, ignored).

subscribe_link(Name, Fun, State) ->
    proc_lib:start(?MODULE, do_subscribe_link, [Name, Fun, State, self()]).


unsubscribe(Pid) ->
    Pid ! unsubscribe,
    misc:wait_for_process(Pid, infinity),

    %% consume exit message in case trap_exit is true
    receive
        %% we expect the process to die normally; if it's not the case then
        %% this should be handled explicitly by parent process;
        {'EXIT', Pid, normal} ->
            ok
    after 0 ->
            ok
    end.

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


init(State) ->
    {ok, State}.


handle_call(_Request, State) ->
    {ok, ok, State}.


handle_event(Event, State = #state{func=Fun, func_state=FS}) ->
    NewState = Fun(Event, FS),
    {ok, State#state{func_state=NewState}}.


handle_info(_Msg, State) ->
    {ok, State}.


terminate(_Args, _State) ->
    ok.


%%
%% Internal functions
%%

%% Function sending a message to a pid
msg_fun(Pid) ->
    fun (Event, ignored) ->
            Pid ! Event,
            ignored
    end.


do_subscribe_link(Name, Fun, State, Parent) ->
    process_flag(trap_exit, true),
    erlang:link(Parent),

    Ref = make_ref(),
    Handler = {?MODULE, Ref},
    ok = gen_event:add_sup_handler(Name, Handler,
                                   #state{func=Fun, func_state=State}),

    proc_lib:init_ack(Parent, self()),

    ExitReason =
        receive
            unsubscribe ->
                normal;
            {'gen_event_EXIT', Handler, Reason} ->
                case Reason =:= normal orelse Reason =:= shutdown of
                    true ->
                        normal;
                    false ->
                        {handler_crashed, Name, Reason}
                end;
            {'EXIT', Parent, Reason} ->
                ?log_debug("Parent process of subscription ~p "
                           "exited with reason ~p", [{Name, Parent}, Reason]),
                normal;
            {'EXIT', Pid, Reason} ->
                ?log_debug("Linked process ~p of subscription ~p "
                           "died unexpectedly with reason ~p",
                           [Pid, {Name, Parent}, Reason]),
                {linked_process_died, Pid, Reason};
            X ->
                ?log_error("Subscription ~p got unexpected message: ~p",
                           [{Name, Parent}, X]),
                unexpected_message
        end,

    (catch gen_event:delete_handler(Name, Handler, unused)),

    exit(ExitReason).

%% Tests
-ifdef(EUNIT).

-define(N, 100).
-define(TABLE, ns_pubsub_test_table).
-define(COUNTER, counter).
-define(EVENTS, ns_pubsub_test_events).

times(0, _F, _A) ->
    ok;
times(N, F, A) ->
    apply(F, A),
    times(N - 1, F, A).

kill_silently_sync(P) ->
    unlink(P),
    exit(P, kill),
    misc:wait_for_process(P, infinity).

create_table() ->
    ?TABLE = ets:new(?TABLE, [public, set, named_table]),
    N = random:uniform(1000000),
    true = ets:insert_new(?TABLE, {?COUNTER, N}),
    ok.

table_updater() ->
    proc_lib:spawn_link(
      fun () ->
              table_updater_loop()
      end).

table_updater_loop() ->
    NewCounter = ets:update_counter(?TABLE, ?COUNTER, 1),
    gen_event:notify(?EVENTS, {counter_updated, NewCounter}),
    table_updater_loop().

setup() ->
    setup(msg_fun(self()), ignored).

setup(Fun, State) ->
    ok = create_table(),
    {ok, EventMgr} = gen_event:start_link({local, ?EVENTS}),
    Updater = table_updater(),
    Subscription = ns_pubsub:subscribe_link(?EVENTS, Fun, State),
    {EventMgr, Updater, Subscription}.

cleanup({EventMgr, Updater, Subscription}) ->
    kill_silently_sync(Updater),
    kill_silently_sync(Subscription),
    kill_silently_sync(EventMgr),

    true = ets:delete(?TABLE),
    flush_all().

wrap(Setup, Cleanup, Body) ->
    R = Setup(),
    try
        Body(R)
    after
        Cleanup(R)
    end.

flush_all() ->
    receive
        _Msg ->
            flush_all()
    after
        0 ->
            ok
    end.

subscribe_test_() ->
    {spawn, [{timeout, 20, ?_test(times(?N, fun test_subscribe/0, []))}]}.

%% test that we don't lose notifications right after subscription
test_subscribe() ->
    wrap(
      fun setup/0, fun cleanup/1,
      fun (_) ->
              [{?COUNTER, Initial}] = ets:lookup(?TABLE, ?COUNTER),

              receive
                  {counter_updated, Value} ->
                      true = (Value =< Initial + 1)
              end
      end).

unsubscribe_test_() ->
    {spawn, [{timeout, 20, ?_test(times(?N, fun test_unsubscribe/0, []))}]}.

%% test that we don't receive new updates after we've unsubscribed
test_unsubscribe() ->
    wrap(
      fun setup/0, fun cleanup/1,
      fun ({_EventMgr, _Updater, Subscription}) ->
              ok = ns_pubsub:unsubscribe(Subscription),
              flush_all(),

              Msg =
                  receive
                      {counter_updated, _} = M ->
                          M
                  after
                      10 ->
                          timeout
                  end,

              ?assertEqual(timeout, Msg)
      end).

kill_test_() ->
    {spawn, [?_test(test_shutdown()),
             ?_test(test_crash()),
             ?_test(test_parent_crash()),
             ?_test(test_event_mgr_crash())]}.

just_fail() ->
    %% NOTE: nif_error is to silence dialyzer
    erlang:nif_error(test_timeout_hit).

test_shutdown() ->
    wrap(
      fun setup/0, fun cleanup/1,
      fun ({EventMgr, Updater, Subscription}) ->
              process_flag(trap_exit, true),

              exit(Updater, kill),
              exit(EventMgr, shutdown),

              receive
                  {'EXIT', Subscription, normal} ->
                      ok
              after
                  1000 ->
                      just_fail()
              end
      end).

test_crash() ->
    Crasher =
        fun (crash, _) ->
                exit(crashed);
            (_, _) ->
                ok
        end,

    wrap(
      fun () -> setup(Crasher, ignored) end, fun cleanup/1,
      fun ({_EventMgr, Updater, Subscription}) ->
              process_flag(trap_exit, true),

              exit(Updater, kill),
              gen_event:notify(?EVENTS, crash),

              receive
                  {'EXIT', Subscription, {handler_crashed, _, {'EXIT', crashed}}} ->
                      ok
              after
                  1000 ->
                      just_fail()
              end
      end).

test_event_mgr_crash() ->
    wrap(
      fun setup/0, fun cleanup/1,
      fun ({EventMgr, Updater, Subscription}) ->
              process_flag(trap_exit, true),

              exit(Updater, kill),
              exit(EventMgr, kill),

              receive
                  {'EXIT', Subscription, {linked_process_died, _, killed}} ->
                      ok
              after
                  1000 ->
                      just_fail()
              end
      end).

test_parent_crash() ->
    process_flag(trap_exit, true),
    Parent = self(),

    Pid =
        proc_lib:spawn(
          fun () ->
                  wrap(
                    fun setup/0, fun cleanup/1,
                    fun ({EventMgr, Updater, _} = R) ->
                            unlink(EventMgr),
                            unlink(Updater),

                            Parent ! ({self(), R}),

                            simply_sleep()
                    end)
          end),

    receive
        {Pid, {EventMgr, Updater, Subscription}} ->
            link(Subscription),
            exit(Pid, crash),

            try
                receive
                    {'EXIT', Subscription, normal} ->
                        ok
                after
                    1000 ->
                        just_fail()
                end
            after
                kill_silently_sync(Updater),
                kill_silently_sync(Subscription),
                kill_silently_sync(EventMgr)
            end
    end.

simply_sleep() ->
    timer:sleep(10000),
    simply_sleep().

-endif.
