epgsql/src/epgsql_sock.erl

%%% Copyright (C) 2009 - Will Glozer.  All rights reserved.
%%% Copyright (C) 2011 - Anton Lebedevich.  All rights reserved.

%%% @doc GenServer holding all connection state (including socket).
%%%
%%% See https://www.postgresql.org/docs/current/static/protocol-flow.html
%%% Commands in PostgreSQL are pipelined: you don't need to wait for reply to
%%% be able to send next command.
%%% Commands are processed (and responses to them are generated) in FIFO order.
%%% eg, if you execute 2 SimpleQuery: #1 and #2, first you get all response
%%% packets for #1 and then all for #2:
%%% > SQuery #1
%%% > SQuery #2
%%% < RowDescription #1
%%% < DataRow #1
%%% < CommandComplete #1
%%% < RowDescription #2
%%% < DataRow #2
%%% < CommandComplete #2
%%%
%%% See epgsql_cmd_connect for network connection and authentication setup


-module(epgsql_sock).

-behavior(gen_server).

-export([start_link/0,
         close/1,
         sync_command/3,
         async_command/4,
         get_parameter/2,
         set_notice_receiver/2,
         get_cmd_status/1,
         cancel/1]).

-export([handle_call/3, handle_cast/2, handle_info/2]).
-export([init/1, code_change/3, terminate/2]).

%% loop callback
-export([on_message/3, on_replication/3]).

%% Comand's APIs
-export([set_net_socket/3, init_replication_state/1, set_attr/3, get_codec/1,
         get_rows/1, get_results/1, notify/2, send/2, send/3, send_multi/2,
         get_parameter_internal/2,
         get_replication_state/1, set_packet_handler/2]).

-export_type([transport/0, pg_sock/0]).

-include("epgsql.hrl").
-include("protocol.hrl").
-include("epgsql_replication.hrl").

-type transport() :: {call, any()}
                   | {cast, pid(), reference()}
                   | {incremental, pid(), reference()}.

-type tcp_socket() :: port(). %gen_tcp:socket() isn't exported prior to erl 18
-type repl_state() :: #repl{}.

-record(state, {mod :: gen_tcp | ssl | undefined,
                sock :: tcp_socket() | ssl:sslsocket() | undefined,
                data = <<>>,
                backend :: {Pid :: integer(), Key :: integer()} | undefined,
                handler = on_message :: on_message | on_replication | undefined,
                codec :: epgsql_binary:codec() | undefined,
                queue = queue:new() :: queue:queue({epgsql_command:command(), any(), transport()}),
                current_cmd :: epgsql_command:command() | undefined,
                current_cmd_state :: any() | undefined,
                current_cmd_transport :: transport() | undefined,
                async :: undefined | atom() | pid(),
                parameters = [] :: [{Key :: binary(), Value :: binary()}],
                rows = [] :: [tuple()],
                results = [],
                sync_required :: boolean() | undefined,
                txstatus :: byte() | undefined,  % $I | $T | $E,
                complete_status :: atom() | {atom(), integer()} | undefined,
                repl :: repl_state() | undefined}).

-opaque pg_sock() :: #state{}.

%% -- client interface --

start_link() ->
    gen_server:start_link(?MODULE, [], []).

close(C) when is_pid(C) ->
    catch gen_server:cast(C, stop),
    ok.

-spec sync_command(epgsql:connection(), epgsql_command:command(), any()) -> any().
sync_command(C, Command, Args) ->
    gen_server:call(C, {command, Command, Args}, infinity).

-spec async_command(epgsql:connection(), cast | incremental,
                    epgsql_command:command(), any()) -> reference().
async_command(C, Transport, Command, Args) ->
    Ref = make_ref(),
    Pid = self(),
    ok = gen_server:cast(C, {{Transport, Pid, Ref}, Command, Args}),
    Ref.

get_parameter(C, Name) ->
    gen_server:call(C, {get_parameter, to_binary(Name)}, infinity).

set_notice_receiver(C, PidOrName) when is_pid(PidOrName);
                                       is_atom(PidOrName) ->
    gen_server:call(C, {set_async_receiver, PidOrName}, infinity).

get_cmd_status(C) ->
    gen_server:call(C, get_cmd_status, infinity).

cancel(S) ->
    gen_server:cast(S, cancel).


%% -- command APIs --

%% send()
%% send_many()

-spec set_net_socket(gen_tcp | ssl, tcp_socket() | ssl:sslsocket(), pg_sock()) -> pg_sock().
set_net_socket(Mod, Socket, State) ->
    State1 = State#state{mod = Mod, sock = Socket},
    setopts(State1, [{active, true}]),
    State1.

-spec init_replication_state(pg_sock()) -> pg_sock().
init_replication_state(State) ->
    State#state{repl = #repl{}}.

-spec set_attr(atom(), any(), pg_sock()) -> pg_sock().
set_attr(backend, {_Pid, _Key} = Backend, State) ->
    State#state{backend = Backend};
set_attr(async, Async, State) ->
    State#state{async = Async};
set_attr(txstatus, Status, State) ->
    State#state{txstatus = Status};
set_attr(codec, Codec, State) ->
    State#state{codec = Codec};
set_attr(sync_required, Value, State) ->
    State#state{sync_required = Value};
set_attr(replication_state, Value, State) ->
    State#state{repl = Value}.

%% XXX: be careful!
-spec set_packet_handler(atom(), pg_sock()) -> pg_sock().
set_packet_handler(Handler, State) ->
    State#state{handler = Handler}.

-spec get_codec(pg_sock()) -> epgsql_binary:codec().
get_codec(#state{codec = Codec}) ->
    Codec.

-spec get_replication_state(pg_sock()) -> repl_state().
get_replication_state(#state{repl = Repl}) ->
    Repl.

-spec get_rows(pg_sock()) -> [tuple()].
get_rows(#state{rows = Rows}) ->
    lists:reverse(Rows).

-spec get_results(pg_sock()) -> [any()].
get_results(#state{results = Results}) ->
    lists:reverse(Results).

-spec get_parameter_internal(binary(), pg_sock()) -> binary() | undefined.
get_parameter_internal(Name, #state{parameters = Parameters}) ->
    case lists:keysearch(Name, 1, Parameters) of
        {value, {Name, Value}} -> Value;
        false                  -> undefined
    end.


%% -- gen_server implementation --

init([]) ->
    {ok, #state{}}.

handle_call({get_parameter, Name}, _From, State) ->
    {reply, {ok, get_parameter_internal(Name, State)}, State};

handle_call({set_async_receiver, PidOrName}, _From, #state{async = Previous} = State) ->
    {reply, {ok, Previous}, State#state{async = PidOrName}};

handle_call(get_cmd_status, _From, #state{complete_status = Status} = State) ->
    {reply, {ok, Status}, State};

handle_call({standby_status_update, FlushedLSN, AppliedLSN}, _From,
            #state{handler = on_replication,
                   repl = #repl{last_received_lsn = ReceivedLSN} = Repl} = State) ->
    send(State, ?COPY_DATA, epgsql_wire:encode_standby_status_update(ReceivedLSN, FlushedLSN, AppliedLSN)),
    Repl1 = Repl#repl{last_flushed_lsn = FlushedLSN,
                      last_applied_lsn = AppliedLSN},
    {reply, ok, State#state{repl = Repl1}};
handle_call({command, Command, Args}, From, State) ->
    Transport = {call, From},
    command_new(Transport, Command, Args, State).

handle_cast({{Method, From, Ref} = Transport, Command, Args}, State)
  when ((Method == cast) or (Method == incremental)),
       is_pid(From),
       is_reference(Ref)  ->
    command_new(Transport, Command, Args, State);

handle_cast(stop, State) ->
    {stop, normal, flush_queue(State, {error, closed})};

handle_cast(cancel, State = #state{backend = {Pid, Key},
                                   sock = TimedOutSock}) ->
    {ok, {Addr, Port}} = case State#state.mod of
                             gen_tcp -> inet:peername(TimedOutSock);
                             ssl -> ssl:peername(TimedOutSock)
                         end,
    SockOpts = [{active, false}, {packet, raw}, binary],
    %% TODO timeout
    {ok, Sock} = gen_tcp:connect(Addr, Port, SockOpts),
    Msg = <<16:?int32, 80877102:?int32, Pid:?int32, Key:?int32>>,
    ok = gen_tcp:send(Sock, Msg),
    gen_tcp:close(Sock),
    {noreply, State}.

handle_info({Closed, Sock}, #state{sock = Sock} = State)
  when Closed == tcp_closed; Closed == ssl_closed ->
    {stop, sock_closed, flush_queue(State#state{sock = undefined}, {error, sock_closed})};

handle_info({Error, Sock, Reason}, #state{sock = Sock} = State)
  when Error == tcp_error; Error == ssl_error ->
    Why = {sock_error, Reason},
    {stop, Why, flush_queue(State, {error, Why})};

handle_info({inet_reply, _, ok}, State) ->
    {noreply, State};

handle_info({inet_reply, _, Status}, State) ->
    {stop, Status, flush_queue(State, {error, Status})};

handle_info({_, Sock, Data2}, #state{data = Data, sock = Sock} = State) ->
    loop(State#state{data = <<Data/binary, Data2/binary>>}).

terminate(_Reason, #state{sock = undefined}) -> ok;
terminate(_Reason, #state{mod = gen_tcp, sock = Sock}) -> gen_tcp:close(Sock);
terminate(_Reason, #state{mod = ssl, sock = Sock}) -> ssl:close(Sock).

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

%% -- internal functions --

-spec command_new(transport(), epgsql_command:command(), any(), pg_sock()) ->
                         Result when
      Result :: {noreply, pg_sock()}
              | {stop, Reason :: any(), pg_sock()}.
command_new(Transport, Command, Args, State) ->
    CmdState = epgsql_command:init(Command, Args),
    command_exec(Transport, Command, CmdState, State).

-spec command_exec(transport(), epgsql_command:command(), any(), pg_sock()) ->
                          Result when
      Result :: {noreply, pg_sock()}
              | {stop, Reason :: any(), pg_sock()}.
command_exec(Transport, Command, _, State = #state{sync_required = true})
  when Command /= epgsql_cmd_sync ->
    {noreply,
     finish(State#state{current_cmd = Command,
                        current_cmd_transport = Transport},
            {error, sync_required})};
command_exec(Transport, Command, CmdState, State) ->
    case epgsql_command:execute(Command, State, CmdState) of
        {ok, State1, CmdState1} ->
            {noreply, command_enqueue(Transport, Command, CmdState1, State1)};
        {stop, StopReason, Response, State1} ->
            reply(Transport, Response, Response),
            {stop, StopReason, State1}
    end.

-spec command_enqueue(transport(), epgsql_command:command(), epgsql_command:state(), pg_sock()) -> pg_sock().
command_enqueue(Transport, Command, CmdState, #state{current_cmd = undefined} = State) ->
    State#state{current_cmd = Command,
                current_cmd_state = CmdState,
                current_cmd_transport = Transport,
                complete_status = undefined};
command_enqueue(Transport, Command, CmdState, #state{queue = Q} = State) ->
    State#state{queue = queue:in({Command, CmdState, Transport}, Q),
                complete_status = undefined}.

-spec command_handle_message(byte(), binary() | epgsql:query_error(), pg_sock()) ->
                                    {noreply, pg_sock()}
                                  | {stop, any(), pg_sock()}.
command_handle_message(Msg, Payload,
                       #state{current_cmd = Command,
                              current_cmd_state = CmdState} = State) ->
    case epgsql_command:handle_message(Command, Msg, Payload, State, CmdState) of
        {add_row, Row, State1, CmdState1} ->
            {noreply, add_row(State1#state{current_cmd_state = CmdState1}, Row)};
        {add_result, Result, Notice, State1, CmdState1} ->
            {noreply,
             add_result(State1#state{current_cmd_state = CmdState1},
                        Notice, Result)};
        {finish, Result, Notice, State1} ->
            {noreply, finish(State1, Notice, Result)};
        {noaction, State1} ->
            {noreply, State1};
        {noaction, State1, CmdState1} ->
            {noreply, State1#state{current_cmd_state = CmdState1}};
        {requeue, State1, CmdState1} ->
            Transport = State1#state.current_cmd_transport,
            command_exec(Transport, Command, CmdState1,
                         State1#state{current_cmd = undefined});
        {stop, Reason, Response, State1} ->
            {stop, Reason, finish(State1, Response)};
        {sync_required, Why} ->
            %% Protocol error. Finish and flush all pending commands.
            {noreply, sync_required(finish(State#state{sync_required = true}, Why))};
        unknown ->
            {stop, {error, {unexpected_message, Msg, Command, CmdState}}, State}
    end.

command_next(#state{current_cmd = PrevCmd,
                    queue = Q} = State) when PrevCmd =/= undefined ->
    case queue:out(Q) of
        {empty, _} ->
            State#state{current_cmd = undefined,
                        current_cmd_state = undefined,
                        current_cmd_transport = undefined,
                        rows = [],
                        results = []};
        {{value, {Command, CmdState, Transport}}, Q1} ->
            State#state{current_cmd = Command,
                        current_cmd_state = CmdState,
                        current_cmd_transport = Transport,
                        queue = Q1,
                        rows = [],
                        results = []}
    end.


setopts(#state{mod = Mod, sock = Sock}, Opts) ->
    case Mod of
        gen_tcp -> inet:setopts(Sock, Opts);
        ssl     -> ssl:setopts(Sock, Opts)
    end.

%% This one only used in connection initiation to send client's
%% `StartupMessage' and `SSLRequest' packets
-spec send(pg_sock(), iodata()) -> ok | {error, any()}.
send(#state{mod = Mod, sock = Sock}, Data) ->
    do_send(Mod, Sock, epgsql_wire:encode_command(Data)).

-spec send(pg_sock(), byte(), iodata()) -> ok | {error, any()}.
send(#state{mod = Mod, sock = Sock}, Type, Data) ->
    do_send(Mod, Sock, epgsql_wire:encode_command(Type, Data)).

-spec send_multi(pg_sock(), [{byte(), iodata()}]) -> ok | {error, any()}.
send_multi(#state{mod = Mod, sock = Sock}, List) ->
    do_send(Mod, Sock, lists:map(fun({Type, Data}) ->
        epgsql_wire:encode_command(Type, Data)
    end, List)).

do_send(gen_tcp, Sock, Bin) ->
    %% Why not gen_tcp:send/2?
    %% See https://github.com/rabbitmq/rabbitmq-common/blob/v3.7.4/src/rabbit_writer.erl#L367-L384
    %% Because of that we also have `handle_info({inet_reply, ...`
    try erlang:port_command(Sock, Bin) of
        true ->
            ok
    catch
        error:_Error ->
            {error, einval}
    end;
do_send(ssl, Sock, Bin) ->
    ssl:send(Sock, Bin).

loop(#state{data = Data, handler = Handler, repl = Repl} = State) ->
    case epgsql_wire:decode_message(Data) of
        {Type, Payload, Tail} ->
            case ?MODULE:Handler(Type, Payload, State#state{data = Tail}) of
                {noreply, State2} ->
                    loop(State2);
                R = {stop, _Reason2, _State2} ->
                    R
            end;
        _ ->
            %% in replication mode send feedback after each batch of messages
            case (Repl =/= undefined) andalso (Repl#repl.feedback_required) of
                true ->
                    #repl{last_received_lsn = LastReceivedLSN,
                          last_flushed_lsn = LastFlushedLSN,
                          last_applied_lsn = LastAppliedLSN} = Repl,
                    send(State, ?COPY_DATA, epgsql_wire:encode_standby_status_update(
                        LastReceivedLSN, LastFlushedLSN, LastAppliedLSN)),
                    {noreply, State#state{repl = Repl#repl{feedback_required = false}}};
                _ ->
                    {noreply, State}
            end
    end.

finish(State, Result) ->
    finish(State, Result, Result).

finish(State = #state{current_cmd_transport = Transport}, Notice, Result) ->
    reply(Transport, Notice, Result),
    command_next(State).

reply({cast, From, Ref}, _, Result) ->
    From ! {self(), Ref, Result};
reply({incremental, From, Ref}, Notice, _) ->
    From ! {self(), Ref, Notice};
reply({call, From}, _, Result) ->
    gen_server:reply(From, Result).

add_result(#state{results = Results, current_cmd_transport = Transport} = State, Notice, Result) ->
    Results2 = case Transport of
                   {incremental, From, Ref} ->
                       From ! {self(), Ref, Notice},
                       Results;
                   _ ->
                       [Result | Results]
               end,
    State#state{rows = [],
                results = Results2}.

add_row(#state{rows = Rows, current_cmd_transport = Transport} = State, Data) ->
    Rows2 = case Transport of
                {incremental, From, Ref} ->
                    From ! {self(), Ref, {data, Data}},
                    Rows;
                _ ->
                    [Data | Rows]
            end,
    State#state{rows = Rows2}.

notify(#state{current_cmd_transport = {incremental, From, Ref}} = State, Notice) ->
    From ! {self(), Ref, Notice},
    State;
notify(State, _) ->
    State.

%% Send asynchronous messages (notice / notification)
notify_async(#state{async = undefined}, _) ->
    false;
notify_async(#state{async = PidOrName}, Msg) ->
    try PidOrName ! {epgsql, self(), Msg} of
        _ -> true
    catch error:badarg ->
            %% no process registered under this name
            false
    end.

sync_required(#state{current_cmd = epgsql_cmd_sync} = State) ->
    State;
sync_required(#state{current_cmd = undefined} = State) ->
    State#state{sync_required = true};
sync_required(State) ->
    sync_required(finish(State, {error, sync_required})).

flush_queue(#state{current_cmd = undefined} = State, _) ->
    State;
flush_queue(State, Error) ->
    flush_queue(finish(State, Error), Error).

to_binary(B) when is_binary(B) -> B;
to_binary(L) when is_list(L)   -> list_to_binary(L).


%% -- backend message handling --

%% CommandComplete
on_message(?COMMAND_COMPLETE = Msg, Bin, State) ->
    Complete = epgsql_wire:decode_complete(Bin),
    command_handle_message(Msg, Bin, State#state{complete_status = Complete});

%% ReadyForQuery
on_message(?READY_FOR_QUERY = Msg, <<Status:8>> = Bin, State) ->
    command_handle_message(Msg, Bin, State#state{txstatus = Status});

%% Error
on_message(?ERROR = Msg, Err, #state{current_cmd = CurrentCmd} = State) ->
    Reason = epgsql_wire:decode_error(Err),
    case CurrentCmd of
        undefined ->
            %% Message generated by server asynchronously
            {stop, {shutdown, Reason}, State};
        _ ->
            command_handle_message(Msg, Reason, State)
    end;

%% NoticeResponse
on_message(?NOTICE, Data, State) ->
    notify_async(State, {notice, epgsql_wire:decode_error(Data)}),
    {noreply, State};

%% ParameterStatus
on_message(?PARAMETER_STATUS, Data, State) ->
    [Name, Value] = epgsql_wire:decode_strings(Data),
    Parameters2 = lists:keystore(Name, 1, State#state.parameters,
                                 {Name, Value}),
    {noreply, State#state{parameters = Parameters2}};

%% NotificationResponse
on_message(?NOTIFICATION, <<Pid:?int32, Strings/binary>>, State) ->
    {Channel1, Payload1} = case epgsql_wire:decode_strings(Strings) of
        [Channel, Payload] -> {Channel, Payload};
        [Channel]          -> {Channel, <<>>}
    end,
    notify_async(State, {notification, Channel1, Pid, Payload1}),
    {noreply, State};

%% ParseComplete
%% ParameterDescription
%% RowDescription
%% NoData
%% BindComplete
%% CloseComplete
%% DataRow
%% PortalSuspended
%% EmptyQueryResponse
%% CopyData
%% CopyBothResponse
on_message(Msg, Payload, State) ->
    command_handle_message(Msg, Payload, State).


%% CopyData for Replication mode
on_replication(?COPY_DATA, <<?PRIMARY_KEEPALIVE_MESSAGE:8, LSN:?int64, _Timestamp:?int64, ReplyRequired:8>>,
               #state{repl = #repl{last_flushed_lsn = LastFlushedLSN,
                                   last_applied_lsn = LastAppliedLSN,
                                   align_lsn = AlignLsn} = Repl} = State) ->
    Repl1 =
        case ReplyRequired of
            1 when AlignLsn ->
                send(State, ?COPY_DATA,
                     epgsql_wire:encode_standby_status_update(LSN, LSN, LSN)),
                Repl#repl{feedback_required = false,
                     last_received_lsn = LSN, last_applied_lsn = LSN, last_flushed_lsn = LSN};
            1 when not AlignLsn ->
                send(State, ?COPY_DATA,
                     epgsql_wire:encode_standby_status_update(LSN, LastFlushedLSN, LastAppliedLSN)),
                Repl#repl{feedback_required = false,
                          last_received_lsn = LSN};
            _ ->
                Repl#repl{feedback_required = true,
                          last_received_lsn = LSN}
        end,
    {noreply, State#state{repl = Repl1}};

%% CopyData for Replication mode
on_replication(?COPY_DATA, <<?X_LOG_DATA, StartLSN:?int64, EndLSN:?int64,
                             _Timestamp:?int64, WALRecord/binary>>,
               #state{repl = Repl} = State) ->
    Repl1 = handle_xlog_data(StartLSN, EndLSN, WALRecord, Repl),
    {noreply, State#state{repl = Repl1}};
on_replication(?ERROR, Err, State) ->
    Reason = epgsql_wire:decode_error(Err),
    {stop, {error, Reason}, State};
on_replication(M, Data, Sock) when M == ?NOTICE;
                                   M == ?NOTIFICATION;
                                   M == ?PARAMETER_STATUS ->
    on_message(M, Data, Sock).


handle_xlog_data(StartLSN, EndLSN, WALRecord, #repl{cbmodule = undefined,
                                                    receiver = Receiver} = Repl) ->
    %% with async messages
    Receiver ! {epgsql, self(), {x_log_data, StartLSN, EndLSN, WALRecord}},
    Repl#repl{feedback_required = true,
              last_received_lsn = EndLSN};
handle_xlog_data(StartLSN, EndLSN, WALRecord,
                 #repl{cbmodule = CbModule, cbstate = CbState, receiver = undefined} = Repl) ->
    %% with callback method
    {ok, LastFlushedLSN, LastAppliedLSN, NewCbState} =
        epgsql:handle_x_log_data(CbModule, StartLSN, EndLSN, WALRecord, CbState),
    Repl#repl{feedback_required = true,
              last_received_lsn = EndLSN,
              last_flushed_lsn = LastFlushedLSN,
              last_applied_lsn = LastAppliedLSN,
              cbstate = NewCbState}.