Fix cancelling timeouted queued queries in transactions for MySQL

lib/dialects/mysql/index.js

@@ -2,6 +2,7 @@
 // -------
 const defer = require('lodash/defer');
 const map = require('lodash/map');
+const remove = require('lodash/remove');
 const { promisify } = require('util');
 const Client = require('../../client');
 
@@ -123,7 +124,7 @@ class Client_MySQL extends Client {
         return;
       }
       const queryOptions = Object.assign({ sql: obj.sql }, obj.options);
-      connection.query(
+      obj.__query = connection.query(
         queryOptions,
         obj.bindings,
         function (err, rows, fields) {
@@ -161,7 +162,17 @@ class Client_MySQL extends Client {
     }
   }
 
-  async cancelQuery(connectionToKill) {
+  _cancelQueuedQuery(connectionToKill, queryToKill) {
+    const removed = remove(connectionToKill._protocol._queue, (query, idx) => (
+      idx > 0 && query === queryToKill.__query
+    ));
+    return removed.length > 0;
+  }
+
+  async cancelQuery(connectionToKill, queryToKill) {
+    if (this._cancelQueuedQuery(connectionToKill, queryToKill)) return;
+
+    connectionToKill.pause();
     const conn = await this.acquireRawConnection();
     try {
       return await this._query(conn, {
@@ -170,6 +181,7 @@ class Client_MySQL extends Client {
         options: {},
       });
     } finally {
+      connectionToKill.resume();
       await this.destroyRawConnection(conn);
       if (conn.__knex__disposed) {
         this.logger.warn(`Connection Error: ${conn.__knex__disposed}`);

lib/dialects/mysql2/index.js

@@ -1,5 +1,6 @@
 // MySQL2 Client
 // -------
+const findIndex = require('lodash/findIndex');
 const Client_MySQL = require('../mysql');
 const Transaction = require('./transaction');
 
@@ -20,6 +21,18 @@ class Client_MySQL2 extends Client_MySQL {
     }
     return true;
   }
+
+  _cancelQueuedQuery(connectionToKill, queryToKill) {
+    const index = findIndex(
+      connectionToKill._commands.toArray(),
+      (cmd) => cmd === queryToKill.__query
+    );
+    if (index >= 0) {
+      connectionToKill._commands.removeOne(index);
+      return true;
+    }
+    return false;
+  }
 }
 
 Object.assign(Client_MySQL2.prototype, {

lib/execution/runner.js

@@ -168,7 +168,7 @@ class Runner {
 
         let cancelQuery;
         if (obj.cancelOnTimeout) {
-          cancelQuery = this.client.cancelQuery(this.connection);
+          cancelQuery = this.client.cancelQuery(this.connection, obj);
         } else {
           // If we don't cancel the query, we need to mark the connection as disposed so that
           // it gets destroyed by the pool and is never used again. If we don't do this and

test/integration/query/additional.js

@@ -978,6 +978,203 @@ module.exports = function (knex) {
         });
     });
 
+    it('.timeout(ms, {cancel: true}) should cancel all parallel slow queries run in a single transaction', async function () {
+      // Only mysql/postgres query cancelling supported for now
+      if (!isMysql(knex) && !isPostgreSQL(knex)) {
+        return this.skip();
+      }
+
+      const testQueries = {
+        [drivers.PostgreSQL]: function () {
+          return knex.raw('SELECT pg_sleep(10)');
+        },
+        [drivers.MySQL]: function () {
+          return knex.raw('SELECT SLEEP(10)');
+        },
+        [drivers.MySQL2]: function () {
+          return knex.raw('SELECT SLEEP(10)');
+        },
+        [drivers.MsSQL]: function () {
+          return knex.raw("WAITFOR DELAY '00:00:10'");
+        },
+        [drivers.Oracle]: function () {
+          return knex.raw('begin dbms_lock.sleep(10); end;');
+        },
+      };
+
+      const driverName = knex.client.driverName;
+      if (!Object.prototype.hasOwnProperty.call(testQueries, driverName)) {
+        throw new Error('Missing test query for dialect: ' + driverName);
+      }
+
+      const query = testQueries[driverName]();
+
+      const getProcessesForDriver = {
+        pg: async () => {
+          const results = await knex.raw('SELECT * from pg_stat_activity');
+          return _.map(_.filter(results.rows, {state: 'active'}), 'query');
+        },
+        mysql: async () => {
+          const results = await knex.raw('SHOW PROCESSLIST');
+          return _.map(results[0], 'Info');
+        },
+        mysql2: async () => {
+          const results = await knex.raw('SHOW PROCESSLIST');
+          return _.map(results[0], 'Info');
+        },
+      };
+
+      if (!Object.prototype.hasOwnProperty.call(getProcessesForDriver, driverName)) {
+        throw new Error('Missing test query for driverName: ' + driverName);
+      }
+
+      const getProcesses = getProcessesForDriver[driverName];
+
+      return knex.transaction(async (trx) => {
+        const promise = Promise.all(_.range(5).map(async () => {
+          await query.timeout(50, { cancel: true }).transacting(trx);
+        }));
+
+        await delay(10);
+        const processesBeforeTimeout = await getProcesses();
+        expect(processesBeforeTimeout).to.include(query.toString());
+
+        await expect(promise).to.eventually.be.rejected.and.to.deep.include({
+          timeout: 50,
+          name: 'KnexTimeoutError',
+          message: 'Defined query timeout of 50ms exceeded when running query.',
+        });
+
+        await delay(10);
+        const processesAfterTimeout = await getProcesses();
+        expect(processesAfterTimeout).to.not.include(query.toString());
+      });
+    });
+
+    it('.timeout(ms, {cancel: true}) should not cancel earlier queries in a transaction before the timeouted query', async function () {
+      // Only mysql/postgres query cancelling supported for now
+      // Postgres automatically aborts the entire transaction if one query is canceled
+      if (!isMysql(knex)) {
+        return this.skip();
+      }
+
+      const getProcessesForDriver = {
+        pg: async () => {
+          const results = await knex.raw('SELECT * from pg_stat_activity');
+          return _.map(_.filter(results.rows, {state: 'active'}), 'query');
+        },
+        mysql: async () => {
+          const results = await knex.raw('SHOW PROCESSLIST');
+          return _.map(results[0], 'Info');
+        },
+        mysql2: async () => {
+          const results = await knex.raw('SHOW PROCESSLIST');
+          return _.map(results[0], 'Info');
+        },
+      };
+
+      const driverName = knex.client.driverName;
+      if (!Object.prototype.hasOwnProperty.call(getProcessesForDriver, driverName)) {
+        throw new Error('Missing test query for driverName: ' + driverName);
+      }
+
+      const getProcesses = getProcessesForDriver[driverName];
+
+      return knex.transaction(async (trx) => {
+        const queryA = trx.raw('SELECT SLEEP(0.1)');
+        const queryB = trx.raw('SELECT SLEEP(0.2)');
+        const promiseA = Promise.all([queryA])
+        const promiseB = Promise.all([queryB.timeout(50, {cancel: true})])
+
+        await delay(10);
+        const processesBeforeTimeout = await getProcesses();
+        expect(processesBeforeTimeout).to.include(queryA.toString());
+        expect(processesBeforeTimeout).to.not.include(queryB.toString());
+
+        await expect(promiseB).to.eventually.be.rejected.and.to.deep.include({
+          timeout: 50,
+          name: 'KnexTimeoutError',
+          message: 'Defined query timeout of 50ms exceeded when running query.',
+        });
+
+        await delay(10);
+        const processesAfterTimeout = await getProcesses();
+        expect(processesAfterTimeout).to.include(queryA.toString());
+        expect(processesAfterTimeout).to.not.include(queryB.toString());
+
+        await expect(promiseA).to.eventually.be.fulfilled.and.to.be.an('array');
+
+        await delay(10);
+        const processesAfterCompletion = await getProcesses();
+        expect(processesAfterCompletion).to.not.include(queryA.toString());
+        expect(processesAfterCompletion).to.not.include(queryB.toString());
+      });
+    });
+
+    it('.timeout(ms, {cancel: true}) should not cancel subsequent queries in a transaction if timeouted query completes while being cancelled', async function () {
+      // Only mysql/postgres query cancelling supported for now
+      // Postgres automatically aborts the entire transaction if one query is canceled
+      if (!isMysql(knex)) {
+        return this.skip();
+      }
+
+      const getProcessesForDriver = {
+        pg: async () => {
+          const results = await knex.raw('SELECT * from pg_stat_activity');
+          return _.map(_.filter(results.rows, {state: 'active'}), 'query');
+        },
+        mysql: async () => {
+          const results = await knex.raw('SHOW PROCESSLIST');
+          return _.map(results[0], 'Info');
+        },
+        mysql2: async () => {
+          const results = await knex.raw('SHOW PROCESSLIST');
+          return _.map(results[0], 'Info');
+        },
+      };
+
+      const driverName = knex.client.driverName;
+      if (!Object.prototype.hasOwnProperty.call(getProcessesForDriver, driverName)) {
+        throw new Error('Missing test query for driverName: ' + driverName);
+      }
+
+      const getProcesses = getProcessesForDriver[driverName];
+
+      return knex.transaction(async (trx) => {
+        const queryA = trx.raw('SELECT SLEEP(1), "A"');
+        const queryB = trx.raw('SELECT SLEEP(1), "B"');
+
+        const promise = Promise.all([
+          queryA.timeout(500, { cancel: true }),
+          queryB,
+        ])
+
+        await delay(100);
+        const processesBeforeTimeout = await getProcesses();
+        expect(processesBeforeTimeout).to.include(queryA.toString());
+        expect(processesBeforeTimeout).to.not.include(queryB.toString());
+
+        const origAcquireRawConnection = trx.client.acquireRawConnection.bind(trx.client);
+        trx.client.acquireRawConnection = async () => {
+          await delay(600);
+          return origAcquireRawConnection();
+        }
+
+        await expect(promise).to.eventually.be.rejected.and.to.deep.include({
+          timeout: 500,
+          name: 'KnexTimeoutError',
+          message: 'Defined query timeout of 500ms exceeded when running query.',
+        });
+
+        trx.client.acquireRawConnection = origAcquireRawConnection;
+
+        await delay(100);
+        const processesAfterTimeout = await getProcesses();
+        expect(processesAfterTimeout).to.not.include(queryA.toString());
+        expect(processesAfterTimeout).to.include(queryB.toString());
+      });
+    });
+
     it('.timeout(ms, {cancel: true}) should cancel slow query even if connection pool is exhausted', async function () {
       // Only mysql/postgres query cancelling supported for now
       if (!isMysql(knex) && !isPostgreSQL(knex)) {