diff --git a/src/main/scala/lc/Main.scala b/src/main/scala/lc/Main.scala index f5a94bd..5bdd849 100644 --- a/src/main/scala/lc/Main.scala +++ b/src/main/scala/lc/Main.scala @@ -25,6 +25,15 @@ object LCFramework { playgroundEnabled = config.playgroundEnabled, corsHeader = config.corsHeader ) + + Runtime.getRuntime.addShutdownHook( new Thread { + override def run(): Unit = { + println("Shutting down gracefully...") + backgroundTask.shutdown() + } + } + ) + server.start() } } diff --git a/src/main/scala/lc/background/taskThread.scala b/src/main/scala/lc/background/taskThread.scala index c7f59ce..431c93e 100644 --- a/src/main/scala/lc/background/taskThread.scala +++ b/src/main/scala/lc/background/taskThread.scala @@ -22,14 +22,18 @@ class BackgroundTask(config: Config, captchaManager: CaptchaManager) { val requiredCountPerCombination = Math.max(1, (config.throttle * 1.01) / allCombinations.size).toInt for (param <- allCombinations) { - val countExisting = captchaManager.getCount(param).getOrElse(0) - val countRequired = requiredCountPerCombination - countExisting - if (countRequired > 0) { - val countCreate = Math.min(1.0 + requiredCountPerCombination/10.0, countRequired).toInt - println(s"Creating $countCreate of $countRequired captchas for $param") + if (!shutdownInProgress) { + val countExisting = captchaManager.getCount(param).getOrElse(0) + val countRequired = requiredCountPerCombination - countExisting + if (countRequired > 0) { + val countCreate = Math.min(1.0 + requiredCountPerCombination/10.0, countRequired).toInt + println(s"Creating $countCreate of $countRequired captchas for $param") - for (i <- 0 until countCreate) { - captchaManager.generateChallenge(param) + for (i <- 0 until countCreate) { + if (!shutdownInProgress) { + captchaManager.generateChallenge(param) + } + } } } } @@ -62,9 +66,22 @@ class BackgroundTask(config: Config, captchaManager: CaptchaManager) { list(HelperFunctions.randomNumber(list.size)) } + private val ex = new ScheduledThreadPoolExecutor(1) + def beginThread(delay: Int): Unit = { - val ex = new ScheduledThreadPoolExecutor(1) ex.scheduleWithFixedDelay(task, 1, delay, TimeUnit.SECONDS) } + @volatile var shutdownInProgress = false + + def shutdown(): Unit = { + println(" Shutting down background task...") + shutdownInProgress = true + ex.shutdown() + println(" Finished Shutting background task") + println(" Shutting down DB...") + Statements.tlStmts.get.shutdown.execute() + println(" Finished shutting down db") + } + } diff --git a/src/main/scala/lc/database/DB.scala b/src/main/scala/lc/database/DB.scala index 207d372..ee6532f 100644 --- a/src/main/scala/lc/database/DB.scala +++ b/src/main/scala/lc/database/DB.scala @@ -3,7 +3,7 @@ package lc.database import java.sql.{Connection, DriverManager, Statement} class DBConn() { - val con: Connection = DriverManager.getConnection("jdbc:h2:./data/H2/captcha2", "sa", "") + val con: Connection = DriverManager.getConnection("jdbc:h2:./data/H2/captcha2;MAX_COMPACT_TIME=8000;DB_CLOSE_ON_EXIT=FALSE", "sa", "") def getStatement(): Statement = { con.createStatement() diff --git a/src/main/scala/lc/database/statements.scala b/src/main/scala/lc/database/statements.scala index 262d9be..bc89fb4 100644 --- a/src/main/scala/lc/database/statements.scala +++ b/src/main/scala/lc/database/statements.scala @@ -120,6 +120,14 @@ class Statements(dbConn: DBConn, maxAttempts: Int) { "SELECT * FROM mapId" ) + val shutdown: PreparedStatement = dbConn.con.prepareStatement( + "SHUTDOWN" + ) + + val shutdownCompact: PreparedStatement = dbConn.con.prepareStatement( + "SHUTDOWN COMPACT" + ) + } object Statements {