package ch.aquint

import com.zaxxer.hikari.HikariConfig
import com.zaxxer.hikari.HikariDataSource
import kotlin.concurrent.thread
import kotlin.system.exitProcess

/**
 * Test case related to https://jira.mariadb.org/browse/CONJ-863 [Aurora SocketTimeout doesn't close Sockets - Zombie Transactions on Read timed out]
 *
 * Test case is structured as following:
 * Precondition:
 * - DB consists of rows 1,3,5,7,9
 * - SocketTimeout is set to 1000ms
 *
 * TestCase:
 * 1. Thread 1 attempts to delete rows 1,3,5,7,9 but runs into a socket timeout because the delete takes longer than anticipated
 * 2. Thread 2 attempts to insert rows 2,4,6,8,10 but fails to do so while the transaction of Thread 1 is running
 * 3. As soon as Thread 3 can successfully insert its rows the test is complete
 *
 * [Official Driver]
 * - Thread 1 has a socket timeout after 1s - but the transaction stays active on the DB because the driver isn't closing the TCP Socket
 *   this causes Thread 2 to be blocked until the previous query completes in the background
 *
 * [Fixed Driver]
 * - Thread 1 has a socket timeout after 1s - upon getting a socket timeout the driver closes the TCP connection
 *   therefore Thread 2 can insert its rows shortly afterwards
 *
 * [Considerations for Fixed Driver]
 * - The fixed driver handles socket timeouts better by freeing up locks earlier and killing the query on the server
 *   this might however have an unwanted impact on current functionality when combining Autocommit + Socket Timeout
 *
 *   [SocketTimeout + Transaction]
 *   - When a transaction is active it makes sense to close the TCP connection ASAP, because the client will never
 *     be able to issue a COMMIT anyway since the connection was thrown away freeing up server resources ASAP
 *
 *   [SocketTimeout + Autocommit]
 *   - When autocommit is ON closing the TCP connection on timeout is a slight change in behaviour because it kills the
 *     query early, since no COMMIT is needed from the Client the DML is probably still applied even after the socket timeout has been reached.
 *     Therefore it could make sense to only close the socket when in a transaction - therefore the potential inTransaction() check.
 *
 * [Proposed Fix]
 * Add this in `AbstractQueryProtocol.handleIoException`
 *
 *     else if (initialException instanceof SocketTimeoutException) {
 *         maxSizeError = false;
 *         try {
 *              // Not sure if inTransaction() check is needed
 *             if (inTransaction()) {
 *                 destroySocket();
 *             }
 *         } catch (Exception e) {
 *             logger.info("Failed to close socket after SocketTimeoutException - " + e.getMessage());
 *         }
 *     }
 */
fun main() {
    // Setup Hikari DataSource
    val hikariConfig = HikariConfig()
    val host = "localhost"
    hikariConfig.jdbcUrl = "jdbc:mariadb://$host:3306/mariadb_zombie_connections?createDatabaseIfNotExist=true&socketTimeout=1000"
    hikariConfig.username = "root"
    hikariConfig.password = "secret"
    val dataSource = HikariDataSource(hikariConfig)

    // Prepare test case
    println("Preparing test ...")
    dataSource.connection.use { connection ->
        connection.prepareStatement("DROP TABLE IF EXISTS test;").execute()
        connection.prepareStatement("CREATE TABLE test(id int auto_increment primary key, message varchar(255) not null)").execute()
        connection.prepareStatement("INSERT INTO test(id, message) VALUES(1, 'test-1')").execute()
        connection.prepareStatement("INSERT INTO test(id, message) VALUES(3, 'test-3')").execute()
        connection.prepareStatement("INSERT INTO test(id, message) VALUES(5, 'test-5')").execute()
        connection.prepareStatement("INSERT INTO test(id, message) VALUES(7, 'test-7')").execute()
        connection.prepareStatement("INSERT INTO test(id, message) VALUES(9, 'test-9')").execute()
    }
    println("Test prepared ...")

    // Start thread to display DB connection statistics
    thread {
        while (true) {
            Thread.sleep(1000)
            dataSource.connection.use { connection ->
                var result = connection.prepareStatement("show status like 'Threads_connected';").executeQuery()
                result.next()
                val connections = result.getString(2)
                result.close()

                result = connection.prepareStatement("SELECT count(*) FROM information_schema.processlist WHERE `DB` = 'mariadb_zombie_connections';").executeQuery()
                result.next()
                val processes = result.getInt(1)
                result.close()

                println("Threads_connected: $connections, Processes: $processes")
            }
        }
    }

    // Attempt to delete rows - for some reason this query is slower than anticipated and runs into socketTimeout
    thread {
        println("Simulating slow deletes to cause a zombie transactions piling up")
        runCatching {
            dataSource.connection.use { connection ->
                connection.prepareStatement("BEGIN").execute()
                // A slow (update) query that runs into the socket timeout but acquires gap locks
                // For example the delete might be slow due to an unexpected high number of relationship rows
                connection.prepareStatement("DELETE FROM test WHERE id IN (1, 3, 5, 7, 9);").execute()
                connection.prepareStatement("SELECT SLEEP(30)").execute() // <--- Here we get a socketTimeout
                connection.prepareStatement("COMMIT").execute() // COMMIT is never reached
            }
        }
    }

    // Sleep to make sure thread 1 is started
    Thread.sleep(1000)

    // Start thread that attempts to insert new unrelated rows (prevented from slow delete with gap locks)
    thread {
        val insertStart = System.currentTimeMillis()
        println("Starting to attempt insert of new row")
        for (i in 1 .. Int.MAX_VALUE) {
            println("Insert attempt $i")
            try {
                dataSource.connection.use { connection ->
                    connection.prepareStatement("BEGIN").execute()
                    connection.prepareStatement("INSERT INTO test(id, message) VALUES(2, 'test-2')").execute()
                    connection.prepareStatement("INSERT INTO test(id, message) VALUES(4, 'test-4')").execute()
                    connection.prepareStatement("INSERT INTO test(id, message) VALUES(6, 'test-6')").execute()
                    connection.prepareStatement("INSERT INTO test(id, message) VALUES(8, 'test-8')").execute()
                    connection.prepareStatement("INSERT INTO test(id, message) VALUES(10, 'test-10')").execute()
                    connection.prepareStatement("COMMIT").execute()
                }
                println("Insert SUCCESS after (${System.currentTimeMillis() - insertStart}ms)!")
                exitProcess(0)
            }
            catch (e: Exception) {
                println("ERROR on insert attempt $i: ${e.message}")
            }
        }
    }
}