Improve query cancellation

Author: porsager

src/connection.js

@@ -75,6 +75,7 @@ function Connection(options, queues = {}, { onopen = noop, onend = noop, onclose
       , connectTimer = timer(connectTimedOut, options.connect_timeout)
 
   let socket = null
+    , cancelMessage
     , result = new Result()
     , incoming = Buffer.alloc(0)
     , needsTypes = options.fetch_types
@@ -139,16 +140,14 @@ function Connection(options, queues = {}, { onopen = noop, onend = noop, onclose
   }
 
   async function cancel({ pid, secret }, resolve, reject) {
-    socket || (socket = await createSocket())
-    if (!socket)
-      return
-
-    socket.removeAllListeners()
-    socket = net.Socket()
-    socket.on('connect', () => socket.write(b().i32(16).i32(80877102).i32(pid).i32(secret).end(16)))
-    socket.once('error', reject)
-    socket.once('close', resolve)
-    connect()
+    try {
+      cancelMessage = b().i32(16).i32(80877102).i32(pid).i32(secret).end(16)
+      await connect()
+      socket.once('error', reject)
+      socket.once('close', resolve)
+    } catch (error) {
+      reject(error)
+    }
   }
 
   function execute(q) {
@@ -955,7 +954,7 @@ function Connection(options, queues = {}, { onopen = noop, onend = noop, onclose
   }
 
   function StartupMessage() {
-    return b().inc(4).i16(3).z(2).str(
+    return cancelMessage || b().inc(4).i16(3).z(2).str(
       Object.entries(Object.assign({
         user,
         database,

tests/index.js

@@ -1774,16 +1774,16 @@ t('Execute', async() => {
 
 t('Cancel running query', async() => {
   const query = sql`select pg_sleep(2)`
-  setTimeout(() => query.cancel(), 50)
+  setTimeout(() => query.cancel(), 200)
   const error = await query.catch(x => x)
   return ['57014', error.code]
 })
 
 t('Cancel piped query', async() => {
   await sql`select 1`
-  const last = sql`select pg_sleep(0.05)`.execute()
+  const last = sql`select pg_sleep(0.2)`.execute()
   const query = sql`select pg_sleep(2) as dig`
-  setTimeout(() => query.cancel(), 10)
+  setTimeout(() => query.cancel(), 100)
   const error = await query.catch(x => x)
   await last
   return ['57014', error.code]