-module(pgo_notifications). -behaviour(gen_statem). -export([start_link/1, start_link/2, listen/2, unlisten/2, report_cb/1]). -export([init/1, callback_mode/0, connected/3, disconnected/3, terminate/3]). -include("pgo_internal.hrl"). -include_lib("kernel/include/logger.hrl"). -type channel() :: unicode:unicode_binary(). -record(data, {conn :: #conn{} | undefined, config :: pgo:pool_config(), backoff :: backoff:backoff(), listeners :: #{reference() => {channel(), pid()}}, listener_channels :: #{channel() => #{reference() => pid()}}}). -spec start_link(pgo:pool_config()) -> gen_statem:start_ret(). start_link(Config) -> gen_statem:start_link(?MODULE, [Config], []). -spec start_link(gen_statem:server_name(), pgo:pool_config()) -> gen_statem:start_ret(). start_link(ServerName, Config) -> gen_statem:start_link(ServerName, ?MODULE, [Config], []). -spec listen(gen_statem:server_ref(), channel()) -> {ok, reference()}. listen(ServerRef, Channel) when is_binary(Channel) -> gen_statem:call(ServerRef, {listen, Channel}). -spec unlisten(gen_statem:server_ref(), reference()) -> ok. unlisten(ServerRef, Ref) -> gen_statem:call(ServerRef, {unlisten, Ref}). init([Config]) -> erlang:process_flag(trap_exit, true), B = backoff:init(1000, 10000), {ok, disconnected, #data{backoff=B, config=Config, listeners=#{}, listener_channels=#{}}, {next_event, internal, connect} }. callback_mode() -> [state_functions, state_enter]. disconnected(enter, _, _Data) -> keep_state_and_data; disconnected(EventType, _, Data=#data{backoff=B, config=Config}) when EventType =:= internal ; EventType =:= timeout ; EventType =:= state_timeout -> try pgo_handler:open(undefined, Config) of {ok, Conn} -> {_, B1} = backoff:succeed(B), {next_state, connected, Data#data{conn=Conn, backoff=B1}}; {error, Error} -> ?LOG_DEBUG("full error connecting to database: ~p", [Error]), ?LOG_INFO(#{at => connecting, reason => Error}, #{report_cb => fun ?MODULE:report_cb/1}), {Backoff, B1} = backoff:fail(B), {next_state, disconnected, Data#data{backoff=B1}, [{state_timeout, Backoff, connect}]} catch throw:Reason -> ?LOG_INFO(#{at => connecting, reason => Reason}, #{report_cb => fun ?MODULE:report_cb/1}), {Backoff, B1} = backoff:fail(B), {next_state, disconnected, Data#data{backoff=B1}, [{state_timeout, Backoff, connect}]} end; disconnected({call, {Pid, _}=From}, {listen, Channel}, Data=#data{listeners=Listeners, listener_channels=ListenerChannels}) -> Ref = erlang:monitor(process, Pid), ListenerChannels1 = maps:update_with(Channel, fun(ChannelListeners) -> ChannelListeners#{Ref => Pid} end, #{Ref => Pid}, ListenerChannels), Listeners1 = Listeners#{Ref => {Channel, Pid}}, {keep_state, Data#data{listeners=Listeners1, listener_channels=ListenerChannels1}, [{reply, From, {eventually, Ref}}, {next_event, internal, connect}]}; disconnected(EventType, EventContent, Data) -> handle_event(EventType, EventContent, Data). connected(enter, _, #data{conn=Conn, listener_channels=ListenerChannels}) -> case map_size(ListenerChannels) > 0 of true -> Channels = maps:keys(ListenerChannels), ListenStatements = [["LISTEN ", Channel, "; "] || Channel <- Channels], Query = ["DO $$BEGIN ", ListenStatements, " END$$"], #{command := do} = pgo_handler:extended_query(Conn, Query, []), keep_state_and_data; false -> keep_state_and_data end; connected({call, {Pid, _}=From}, {listen, Channel}, Data=#data{conn=Conn, listeners=Listeners, listener_channels=ListenerChannels}) -> Ref = erlang:monitor(process, Pid), ListenerChannels1 = maps:update_with(Channel, fun(Channels) -> Channels#{Ref => Pid} end, #{Ref => Pid}, ListenerChannels), Listeners1 = Listeners#{Ref => {Channel, Pid}}, case erlang:map_size(maps:get(Channel, ListenerChannels1)) of 1 -> %% first time listening on this channel case pgo_handler:extended_query(Conn, ["LISTEN ", Channel], []) of #{command := listen} -> {keep_state, Data#data{listeners=Listeners1, listener_channels=ListenerChannels1}, [{reply, From, {ok, Ref}}]}; _ -> {keep_state_and_data, [{reply, From, error}]} end; _ -> {keep_state, Data#data{listeners=Listeners1, listener_channels=ListenerChannels1}, [{reply, From, {ok, Ref}}]} end; connected(EventType, EventContent, Data) -> handle_event(EventType, EventContent, Data). handle_event(cast, {set_parameter, Name, Value}, Data=#data{conn=Conn=#conn{parameters=Parameters}}) -> {keep_state, Data#data{conn=Conn#conn{parameters=Parameters#{Name => Value}}}}; handle_event({call, From}, {unlisten, Ref}, Data=#data{conn=Conn, listeners=Listeners, listener_channels=ListenerChannels}) -> {Listeners1, ListenerChannels1} = unlisten(Ref, Listeners, ListenerChannels, Conn), {keep_state, Data#data{listeners=Listeners1, listener_channels=ListenerChannels1}, [{reply, From, ok}]}; handle_event(info, {'DOWN', Ref, process, _, _}, Data=#data{conn=Conn, listeners=Listeners, listener_channels=ListenerChannels}) -> {Listeners1, ListenerChannels1} = unlisten(Ref, Listeners, ListenerChannels, Conn), {keep_state, Data#data{listeners=Listeners1, listener_channels=ListenerChannels1}}; handle_event(info, {Tag, Socket, Binary}, Data=#data{conn=Conn=#conn{socket=Socket}, listener_channels=ListenerChannels}) when Tag =:= tcp orelse Tag =:= ssl -> Conn = pgo_handler:process_active_data(Binary, Conn, fun(Notification) -> notify(ListenerChannels, Notification) end), _ = inet:setopts(Socket, [{active, once}]), {keep_state, Data#data{conn=Conn}}; handle_event(info, {Tag, _, _Binary}, _Data) when Tag =:= tcp orelse Tag =:= ssl -> %% ignore message from an old socket %% this can happen for instance if a message gets into the mailbox before the %% connection is closed and reopen keep_state_and_data; handle_event(info, {'EXIT', Socket, _Reason}, Data=#data{conn=#conn{socket=Socket}}) -> %% socket died, go to disconnected state close_and_reopen(Data); %% ignore `EXIT' for a different Socket -- means it is an old message handle_event(info, {'EXIT', _, _Reason}, _Data) -> keep_state_and_data; %% nothing to do for `ssl_closed' -- it should be handled by the `EXIT' handling handle_event(info, {ssl_closed, _}, _Data) -> keep_state_and_data; handle_event(_, _, _) -> keep_state_and_data. terminate(_Reason, _, #data{conn=undefined}) -> ok; terminate(_Reason, _, #data{conn=Conn}) -> pgo_handler:close(Conn), ok. %% close_and_reopen(Data=#data{conn=Conn}) -> pgo_handler:close(Conn), {next_state, disconnected, Data#data{conn=undefined}, [{next_event, internal, connect}]}. report_cb(#{at := connecting, reason := {pgo_error, #{message := Message}}}) -> {"error connecting to database: ~s", [Message]}; report_cb(#{at := connecting, reason := Reason}) -> {"unknown error when connecting to database: ~p", [Reason]}. unlisten(Ref, Listeners, ListenerChannels, Conn) -> erlang:demonitor(Ref, [flush]), case maps:find(Ref, Listeners) of {ok, {Channel, _Pid}} -> Listeners1 = maps:remove(Ref, Listeners), PerChannel = maps:remove(Ref, maps:get(Channel, ListenerChannels)), case erlang:map_size(PerChannel) of 0 -> ListenerChannels1 =maps:remove(Channel, ListenerChannels), #{command := unlisten} = pgo_handler:extended_query(Conn, ["UNLISTEN ", Channel], []), {Listeners1, ListenerChannels1}; _ -> {Listeners1, ListenerChannels#{Channel => PerChannel}} end; error -> {Listeners, ListenerChannels} end. notify(ListenerChannels, #notification_response{channel=Channel, payload=Payload}) -> case maps:find(Channel, ListenerChannels) of {ok, Listeners} -> maps:foreach(fun(Ref, Pid) -> Pid ! {notification, self(), Ref, Channel, Payload} end, Listeners); error -> ok end.