const mysql = require('mysql'); // npm install mysql
const util = require('util');

/**
 * Tweak the following globals to fit your environment
 * ###################################################
 */
const HOST = '127.0.0.1';
const PORT = 3306;
const USER = 'root';
const PASSWORD = '';

// Specify which database and table to work with.
// Note: this database will be dropped at the end of this script
const DATABASE = 'test';
const TABLE = 'tbl';

// The number of workers to run
const NUM_WORKERS = 20;

// Run the workload for this many seconds
const WORKLOAD_TIME = 10;

// Batch size to use
const BATCH_SIZE = 5000;

/**
 * Internal code starts here
 * #########################
 */

let isDone = false;

// await-able setTimeout()
function timeout(ms) {
  return new Promise(resolve => setTimeout(resolve, ms));
}

// Pre-generate the insert query
const _batch = new Array(BATCH_SIZE).fill().map(_ => '()').join(',');
const insertQuery = `INSERT INTO ${TABLE} VALUES ${_batch}`;;

function getConnection(dbName) {
  return new Promise(function(resolve, reject) {

    const conn = mysql.createConnection({
      host: HOST,
      port: PORT,
      user: USER,
      password: PASSWORD,
      database: dbName
    });

    conn.connect(err => {
      if (err) {
        reject(err);
      } else {
        conn.query = util.promisify(conn.query);
        resolve(conn);
      }
    });

  });
};

async function setupTestDb() {
  const conn = await getConnection('information_schema');
  await conn.query(`CREATE DATABASE IF NOT EXISTS ${DATABASE}`);
  await conn.query(`USE ${DATABASE}`);
  await conn.query(`CREATE TABLE IF NOT EXISTS ${TABLE} (id INT AUTO_INCREMENT PRIMARY KEY)`);
}

async function insertWorker() {
  const conn = await getConnection(DATABASE);
  while (true) {
    // await will process.nextTick()
    await conn.query(insertQuery);
    if (isDone) {
      break;
    }
  }
}

async function warmup() {
  console.log('Warming up workload');
  const conn = await getConnection(DATABASE)
  await conn.query(insertQuery); // FRAGILE: included in count, not included in time
}

async function doBenchmark() {
  console.log(`Launching ${NUM_WORKERS} workers for ${WORKLOAD_TIME} sec`);

  const workers = [];
  for (let i = 0; i < NUM_WORKERS; ++i) {
    workers.push(insertWorker());
  }
  console.log(`${workers.length} workers running...`);

  await timeout(WORKLOAD_TIME * 1000);

  console.log('Stopping workload');
  isDone = true;

  await Promise.all(workers);
}

async function printStats() {
  const conn = await getConnection(DATABASE);
  const rows = await conn.query(`SELECT COUNT(*) AS count FROM ${TABLE}`);
  const count = rows[0].count;
  console.log(`${count} rows inserted using ${NUM_WORKERS} workers`);
  console.log(`${count / WORKLOAD_TIME} rows per second`);
}

async function cleanupTestDb() {
  console.log('Cleaning up');
  const conn = await getConnection('information_schema');
  await conn.query(`DROP DATABASE ${DATABASE}`);
}

async function main() {
  try {
    await setupTestDb();
    await warmup();
    await doBenchmark();
    await printStats();
    await cleanupTestDb();
  } catch (err) {
    console.error('ERROR', err);
    try {
      await cleanupTestDb();
    } catch (err2) {
      console.error(err2);
    }
    process.exit(1);
  }
  process.exit(0); // releases all connections
}

main();