Merge pull request #223 from enidgjoleka/fix-connection-timeout

Do not re-use the same connection timeout during different steps of establishing the connection

.gitignore

@@ -4,3 +4,5 @@ datadir/
 doc/*
 !doc/overview.edoc
 !doc/*.md
+*.idea/*
+*.iml

README.md

@@ -97,8 +97,8 @@ Only `host` and `username` are mandatory, but most likely you would need `databa
 - `password` - DB user password. It might be provided as string / binary or as a fun that returns
    string / binary. Internally, plain password is wrapped to anonymous fun before it is sent to connection
    process, so, if `connect` command crashes, plain password will not appear in crash logs.
-- `{timeout, TimeoutMs}` parameter will trigger an `{error, timeout}` result when the
-   socket fails to connect within `TimeoutMs` milliseconds.
+- `timeout` parameter will trigger an `{error, timeout}` result when the
+   socket fails to connect within provided milliseconds.
 - `ssl` if set to `true`, perform an attempt to connect in ssl mode, but continue unencrypted
   if encryption isn't supported by server. if set to `required` connection will fail if encryption
   is not available.

src/commands/epgsql_cmd_connect.erl

@@ -49,19 +49,20 @@ init(#{host := _, username := _} = Opts) ->
 
 execute(PgSock, #connect{opts = #{host := Host} = Opts, stage = connect} = State) ->
     Timeout = maps:get(timeout, Opts, 5000),
+    Deadline = deadline(Timeout),
     Port = maps:get(port, Opts, 5432),
     SockOpts = [{active, false}, {packet, raw}, binary, {nodelay, true}, {keepalive, true}],
     case gen_tcp:connect(Host, Port, SockOpts, Timeout) of
         {ok, Sock} ->
-            client_handshake(Sock, PgSock, State);
+            client_handshake(Sock, PgSock, State, Deadline);
         {error, Reason} = Error ->
             {stop, Reason, Error, PgSock}
     end;
 execute(PgSock, #connect{stage = auth, auth_send = {PacketId, Data}} = St) ->
-    epgsql_sock:send(PgSock, PacketId, Data),
+    ok = epgsql_sock:send(PgSock, PacketId, Data),
     {ok, PgSock, St#connect{auth_send = undefined}}.
 
-client_handshake(Sock, PgSock, #connect{opts = #{username := Username} = Opts} = State) ->
+client_handshake(Sock, PgSock, #connect{opts = #{username := Username} = Opts} = State, Deadline) ->
     %% Increase the buffer size.  Following the recommendation in the inet man page:
     %%
     %%    It is recommended to have val(buffer) >=
@@ -71,7 +72,7 @@ client_handshake(Sock, PgSock, #connect{opts = #{username := Username} = Opts} =
         inet:getopts(Sock, [recbuf, sndbuf]),
     inet:setopts(Sock, [{buffer, max(RecBufSize, SndBufSize)}]),
 
-    case maybe_ssl(Sock, maps:get(ssl, Opts, false), Opts, PgSock) of
+    case maybe_ssl(Sock, maps:get(ssl, Opts, false), Opts, PgSock, Deadline) of
         {error, Reason} ->
             {stop, Reason, {error, Reason}, PgSock};
         PgSock1 ->
@@ -88,8 +89,7 @@ client_handshake(Sock, PgSock, #connect{opts = #{username := Username} = Opts} =
                          epgsql_sock:init_replication_state(PgSock1)};
                     _ -> {Opts3, PgSock1}
                 end,
-
-            epgsql_sock:send(PgSock2, [<<196608:?int32>>, Opts4, 0]),
+            ok = epgsql_sock:send(PgSock2, [<<196608:?int32>>, Opts4, 0]),
             PgSock3 = case Opts of
                           #{async := Async} ->
                               epgsql_sock:set_attr(async, Async, PgSock2);
@@ -119,15 +119,15 @@ hide_password(PasswordFun) when is_function(PasswordFun, 0) ->
     PasswordFun.
 
 
-maybe_ssl(S, false, _, PgSock) ->
+maybe_ssl(S, false, _, PgSock, _Deadline) ->
     epgsql_sock:set_net_socket(gen_tcp, S, PgSock);
-maybe_ssl(S, Flag, Opts, PgSock) ->
+maybe_ssl(S, Flag, Opts, PgSock, Deadline) ->
     ok = gen_tcp:send(S, <<8:?int32, 80877103:?int32>>),
-    Timeout = maps:get(timeout, Opts, 5000),
-    {ok, <<Code>>} = gen_tcp:recv(S, 1, Timeout),
-    case Code of
-        $S  ->
+    Timeout0 = timeout(Deadline),
+    case gen_tcp:recv(S, 1, Timeout0) of
+        {ok, <<$S>>}  ->
             SslOpts = maps:get(ssl_opts, Opts, []),
+            Timeout = timeout(Deadline),
             case ssl:connect(S, SslOpts, Timeout) of
                 {ok, S2}        ->
                     epgsql_sock:set_net_socket(ssl, S2, PgSock);
@@ -135,13 +135,15 @@ maybe_ssl(S, Flag, Opts, PgSock) ->
                     Err = {ssl_negotiation_failed, Reason},
                     {error, Err}
             end;
-        $N ->
+        {ok, <<$N>>} ->
             case Flag of
                 true ->
                     epgsql_sock:set_net_socket(gen_tcp, S, PgSock);
                 required ->
                     {error, ssl_not_available}
-            end
+            end;
+        {error, Reason} ->
+            {error, Reason}
     end.
 
 %% Auth sub-protocol
@@ -277,3 +279,9 @@ hex(Bin) ->
                (N) when N < 16 -> $W + N
             end,
     <<<<(HChar(H)), (HChar(L))>> || <<H:4, L:4>> <= Bin>>.
+
+deadline(Timeout) ->
+    erlang:monotonic_time(milli_seconds) + Timeout.
+
+timeout(Deadline) ->
+    erlang:max(0, Deadline - erlang:monotonic_time(milli_seconds)).