Let cancel_query! cancel a statement after an interrupt closed the socket

docs/src/manual.md

@@ -278,6 +278,12 @@ Postgres.cancel_query!(conn)
 
 The running query raises a `Postgres.Error` with PostgreSQL SQLSTATE `57014`.
 
+The request carries only the backend pid and cancel key, so it also works on a
+connection that an `InterruptException` closed mid-query. Its statement keeps
+running on the server until it finishes or is cancelled. Call `cancel_query!`
+before you use the connection again: with `reconnect=true`, the next use opens a
+new session, and a later cancel goes to that session instead.
+
 ## LISTEN And NOTIFY
 
 Use one connection to listen and another to notify.

src/Postgres.jl

@@ -829,28 +829,26 @@ the cancel key it carries is a credential: if `conn` itself is on TLS or
 GSSAPI encryption, the cancel connection requires the same. Throws a `PostgresInterfaceError` if the cancel
 request could not be delivered (rather than failing silently, which would
 leave the query running).
+
+The request carries only the backend pid and cancel key, which `conn` keeps
+after its socket closes. So it also cancels the statement of a connection an
+`InterruptException` closed mid-query: the socket is closed because the
+stream position is unknown, but the statement keeps running on the server
+until it finishes or is cancelled. Cancel before using `conn` again: with
+`reconnect=true`, the next use opens a new session, and a later cancel goes to
+that session instead.
 """
 function cancel_query!(conn::Connection)
+    # Read without the lock: cancel_query! is called precisely when another
+    # task holds it running the query being cancelled, and the socket state
+    # does not matter, since the request goes over its own connection.
     host = conn.host
     port = conn.port
     pid = conn.pid
     skey = conn.skey
     debug = conn.debug
-    # Checked outside the lock deliberately: cancel_query! is called precisely
-    # when another task holds it running the query being cancelled, so a
-    # trylock-guarded check would be skipped in the case that matters. Reading
-    # the socket field unlocked matches how host/pid/skey are read below.
     sslmode = cancel_sslmode(conn.socket.transport isa Reseau.TLS.Conn, conn.sslmode)
     gssencmode = cancel_gssencmode(conn.socket.transport isa API.GSSConn, conn.gssencmode)
-    if trylock(conn.lock)
-        try
-            !isopen(conn.socket) && throw(PostgresInterfaceError("cannot cancel query: connection not open"))
-            pid = conn.pid
-            skey = conn.skey
-        finally
-            unlock(conn.lock)
-        end
-    end
     ok = API.cancel_request(host, port, pid, skey, debug, sslmode, conn.sslrootcert, conn.sslcert, conn.sslkey, conn.sslcapath, conn.sslservername, conn.connect_timeout, gssencmode, conn.krbsrvname, conn.gssdelegation, conn.style)
     ok || throw(PostgresInterfaceError("failed to deliver the cancel request to $(host):$(port)"))
     return conn

test/execute_fake_server.jl

@@ -88,6 +88,39 @@ function test_execute_fake_server()
         end
     end
 
+    @testset "cancel after an interrupt closed the socket" begin
+        # An InterruptException that lands mid-statement closes the socket (the
+        # stream position is unknown) and leaves the statement running on the
+        # server. The cancel needs only the pid and key the connection holds.
+        segment_read = Base.Event()
+        cancels = Tuple{Int32, Int32, Int32}[]
+        handler = function (sock, n)
+            if n == 1
+                serve_idle(sock -> begin
+                    read_unnamed_segment(sock)
+                    notify(segment_read)
+                    drain(sock)
+                end)(sock, n)
+            else
+                code, body = read_request(sock)
+                push!(cancels, (code, rd32(body[1:4]), rd32(body[5:8])))
+            end
+        end
+        with_fake_server(handler) do port, accepted
+            conn = execute_connection(port)
+            task = @async connect_err(() -> DBInterface.execute(conn, "SELECT pg_sleep(60)"))
+            wait(segment_read)
+            schedule(task, InterruptException(); error=true)
+            @test fetch(task) isa InterruptException
+            @test !isopen(conn)
+            Postgres.cancel_query!(conn)
+            @test timedwait(() -> !isempty(cancels), 5.0) === :ok
+            # the pid and key the fake server sent in BackendKeyData
+            @test cancels == [(CANCEL_REQUEST, Int32(4242), Int32(777))]
+            @test accepted[] == 2
+        end
+    end
+
     @testset "a server error before ParseComplete keeps the connection" begin
         with_fake_server(serve_idle(sock -> begin
             read_unnamed_segment(sock)

test/runtests.jl

@@ -2859,6 +2859,25 @@ include("trim_compile_tests.jl")
                     @test result.code == "57014"
                     DBInterface.close!(cancel_conn)
 
+                    # an interrupt mid-statement closes the socket but leaves the
+                    # statement running; the cancel still reaches it with the pid
+                    # and key the connection holds
+                    interrupted = DBInterface.connect(Postgres.Connection, cfg.host, cfg.user, cfg.password; dbname=cfg.dbname, port=cfg.port)
+                    pid = interrupted.pid
+                    task = @async connect_err(() -> DBInterface.execute(interrupted, "SELECT pg_sleep(30)"))
+                    started = timedwait(5.0; pollint=0.05) do
+                        only(DBInterface.execute(conn, "SELECT count(*) AS n FROM pg_stat_activity WHERE pid = $pid AND state = 'active'")).n == 1
+                    end
+                    @test started === :ok
+                    schedule(task, InterruptException(); error=true)
+                    @test fetch(task) isa InterruptException
+                    @test !isopen(interrupted)
+                    Postgres.cancel_query!(interrupted)
+                    stopped = timedwait(10.0; pollint=0.05) do
+                        only(DBInterface.execute(conn, "SELECT count(*) AS n FROM pg_stat_activity WHERE pid = $pid AND state = 'active'")).n == 0
+                    end
+                    @test stopped === :ok
+
                     # the cancel key must never go out in the clear when the
                     # connection it cancels required TLS; the refusal is a
                     # thrown error, not a silent no-op (this server has no SSL,