Add tests for verifying the proposed solution and fix discovered issues

src/commands/epgsql_cmd_connect.erl

@@ -49,21 +49,21 @@ init(#{host := _, username := _} = Opts) ->
 
 execute(PgSock, #connect{opts = #{username := Username} = Opts, stage = connect} = State) ->
     SockOpts = [{active, false}, {packet, raw}, binary, {nodelay, true}, {keepalive, true}],
-    epgsql_sock:set_attr(connect_opts, Opts, PgSock),
+    PgSock1 = epgsql_sock:set_attr(connect_opts, Opts, PgSock),
     case open_socket(SockOpts, Opts) of
         {ok, Mode, Sock} ->
-            PgSock1 = epgsql_sock:set_net_socket(Mode, Sock, PgSock),
+            PgSock2 = epgsql_sock:set_net_socket(Mode, Sock, PgSock1),
             Opts2 = ["user", 0, Username, 0],
             Opts3 = case maps:find(database, Opts) of
                         error -> Opts2;
                         {ok, Database}  -> [Opts2 | ["database", 0, Database, 0]]
                     end,
-           {Opts4, PgSock2} =
+           {Opts4, PgSock3} =
                case Opts of
                    #{replication := Replication}  ->
                        {[Opts3 | ["replication", 0, Replication, 0]],
-                        epgsql_sock:init_replication_state(PgSock1)};
-                   _ -> {Opts3, PgSock1}
+                        epgsql_sock:init_replication_state(PgSock2)};
+                   _ -> {Opts3, PgSock2}
                end,
             Opts5 = case Opts of
                         #{application_name := ApplicationName}  ->
@@ -72,12 +72,12 @@ execute(PgSock, #connect{opts = #{username := Username} = Opts, stage = connect}
                             Opts4
                     end,
            ok = epgsql_sock:send(PgSock2, [<<196608:?int32>>, Opts5, 0]),
-           PgSock3 = case Opts of
+           PgSock4 = case Opts of
                          #{async := Async} ->
-                             epgsql_sock:set_attr(async, Async, PgSock2);
-                         _ -> PgSock2
+                             epgsql_sock:set_attr(async, Async, PgSock3);
+                         _ -> PgSock3
                      end,
-           {ok, PgSock3, State#connect{stage = maybe_auth}};
+           {ok, PgSock4, State#connect{stage = maybe_auth}};
         {error, Reason} = Error ->
             {stop, Reason, Error, PgSock}
     end;
@@ -95,7 +95,7 @@ open_socket(SockOpts, #{host := Host} = ConnectOpts) ->
        {error, _Reason} = Error ->
            Error
     end.
-  
+
 client_handshake(Sock, ConnectOpts, Deadline) ->
     %% Increase the buffer size.  Following the recommendation in the inet man page:
     %%

src/epgsql_sock.erl

@@ -228,18 +228,20 @@ 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,
+                                   connect_opts = ConnectOpts,
+                                   mod = Mode}) ->
     SockOpts = [{active, false}, {packet, raw}, binary],
-    %% TODO timeout
-    %% TODO DO NOT use gen_tcp
-    {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),
+    case epgsql_cmd_connect:open_socket(SockOpts, ConnectOpts) of
+      {ok, Mode, Sock} when Mode == gen_tcp ->
+          ok = gen_tcp:send(Sock, Msg),
+          gen_tcp:close(Sock);
+      {ok, Mode, Sock} when Mode == ssl ->
+          ok = ssl:send(Sock, Msg),
+          ssl:close(Sock);
+      {error, _Reason} ->
+          noop
+    end,
     {noreply, State}.
 
 handle_info({Closed, Sock}, #state{sock = Sock} = State)

test/epgsql_SUITE.erl

@@ -45,6 +45,8 @@ groups() ->
             connect_to_invalid_database,
             connect_with_other_error,
             connect_with_ssl,
+            cancel_query_for_connection_with_ssl,
+            cancel_query_for_connection_with_gen_tcp,
             connect_with_client_cert,
             connect_with_invalid_client_cert,
             connect_to_closed_port,
@@ -284,6 +286,63 @@ connect_with_ssl(Config) ->
         "epgsql_test",
         [{ssl, true}]).
 
+cancel_query_for_connection_with_ssl(Config) ->
+    Module = ?config(module, Config),
+    {Host, Port} = epgsql_ct:connection_data(Config),
+    Module = ?config(module, Config),
+    Args2 = [{port, Port}, {database, "epgsql_test_db1"} | [ {ssl, true}, {timeout, 1000} ]],
+    {ok, C} = Module:connect(Host, "epgsql_test", Args2),
+    ?assertMatch({ok, _Cols, [{true}]},
+                Module:equery(C, "select ssl_is_used()")),
+    process_flag(trap_exit, true),
+    Self = self(),
+    spawn_link(fun() ->
+                   ?assertMatch({error, #error{codename = query_canceled}},
+                                Module:equery(C, "SELECT pg_sleep(5)")),
+                   Self ! done
+               end),
+    %% this will never match but introduces 1 second latency needed
+    %% for the test not to be flaky
+    receive none ->
+        noop
+    after 1000 ->
+        epgsql:cancel(C),
+        receive done ->
+            ?assert(true)
+        after 5000 ->
+            epgsql:close(C),
+            ?assert(false)
+        end
+    end.
+
+cancel_query_for_connection_with_gen_tcp(Config) ->
+    Module = ?config(module, Config),
+    {Host, Port} = epgsql_ct:connection_data(Config),
+    Module = ?config(module, Config),
+    Args2 = [{port, Port}, {database, "epgsql_test_db1"} | [ {timeout, 1000} ]],
+    {ok, C} = Module:connect(Host, "epgsql_test", Args2),
+  
+    process_flag(trap_exit, true),
+    Self = self(),
+    spawn_link(fun() ->
+                   ?assertMatch({error, #error{codename = query_canceled}},
+                                Module:equery(C, "SELECT pg_sleep(5)")),
+                   Self ! done
+               end),
+    %% this is will never match but it introduces the 1 second latency needed
+    %% for the test not to be flaky
+    receive none ->
+        noop
+    after 1000 ->
+        epgsql:cancel(C),
+        receive done ->
+            ?assert(true)
+        after 5000 ->
+            epgsql:close(C),
+            ?assert(false)
+        end
+    end.
+
 connect_with_client_cert(Config) ->
     Module = ?config(module, Config),
     Dir = filename:join(code:lib_dir(epgsql), ?TEST_DATA_DIR),