Add timeout support for queries

fixes on timeout scheduler fixes on timeout scheduler fixes on timeout scheduler styling issues Improve unit test

Author: lifey

db-async-common/src/main/scala/com/github/mauricio/async/db/Configuration.scala

@@ -54,4 +54,5 @@ case class Configuration(username: String,
  maximumMessageSize: Int = 16777216,
  allocator: AbstractByteBufAllocator = PooledByteBufAllocator.DEFAULT,
  connectTimeout: Duration = 5.seconds,
- testTimeout: Duration = 5.seconds)
+ testTimeout: Duration = 5.seconds,
+ queryTimeout: Duration = Duration.Inf)

db-async-common/src/main/scala/com/github/mauricio/async/db/exceptions/ConnectionTimeoutedException.scala

@@ -0,0 +1,6 @@
+package com.github.mauricio.async.db.exceptions
+
+import com.github.mauricio.async.db.Connection
+
+class ConnectionTimeoutedException( val connection : Connection )
+ extends DatabaseException( "The connection %s has a timeouted query and is being closed".format(connection) )

db-async-common/src/main/scala/com/github/mauricio/async/db/pool/TimeoutScheduler.scala

@@ -0,0 +1,37 @@
+package com.github.mauricio.async.db.pool
+
+import java.util.concurrent.atomic.AtomicBoolean
+import java.util.concurrent.{TimeUnit, TimeoutException, ScheduledFuture}
+import com.github.mauricio.async.db.util.NettyUtils
+import scala.concurrent.{ExecutionContext, Promise}
+import scala.concurrent.duration.Duration
+
+trait TimeoutScheduler {
+ implicit val internalPool: ExecutionContext
+ def onTimeout // implementors should decide here what they want to do when a timeout occur
+ private var isTimeoutedBool = new AtomicBoolean(false);
+ def isTimeouted = isTimeoutedBool.get // We need this property as isClosed takes time to complete and
+ // we don't want the connection to be used again.
+
+ def addTimeout[A](promise: Promise[A], duration: Duration) : Option[ScheduledFuture[_]] = {
+ if (duration != Duration.Inf) {
+ val scheduledFuture = schedule(
+ {
+ if (promise.tryFailure(new TimeoutException(s"Operation is timeouted after it took too long to return (${duration})"))) {
+ isTimeoutedBool.set(true)
+ onTimeout
+ }
+ },
+ duration)
+ promise.future.onComplete(x => scheduledFuture.cancel(false))
+
+ return Some(scheduledFuture)
+ }
+ return None
+ }
+
+ def schedule(block: => Unit, duration: Duration) : ScheduledFuture[_] =
+ NettyUtils.DefaultEventLoopGroup.schedule(new Runnable {
+ override def run(): Unit = block
+ }, duration.toMillis, TimeUnit.MILLISECONDS)
+}

db-async-common/src/test/scala/com/github/mauricio/async/db/pool/DummyTimeoutScheduler.scala

@@ -0,0 +1,29 @@
+/*
+ * Copyright 2013 Maurício Linhares
+ *
+ * Maurício Linhares licenses this file to you under the Apache License,
+ * version 2.0 (the "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at:
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+
+package com.github.mauricio.async.db.pool
+
+import java.util.concurrent.atomic.AtomicInteger
+import com.github.mauricio.async.db.util.ExecutorServiceUtils
+/**
+ * Implementation of TimeoutScheduler used for testing
+ */
+class DummyTimeoutScheduler extends TimeoutScheduler {
+ implicit val internalPool = ExecutorServiceUtils.CachedExecutionContext
+ private val timeOuts = new AtomicInteger
+ override def onTimeout = timeOuts.incrementAndGet
+ def timeoutCount = timeOuts.get()
+}

db-async-common/src/test/scala/com/github/mauricio/async/db/pool/TimeoutSchedulerSpec.scala

@@ -0,0 +1,71 @@
+/*
+ * Copyright 2013 Maurício Linhares
+ *
+ * Maurício Linhares licenses this file to you under the Apache License,
+ * version 2.0 (the "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at:
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+package com.github.mauricio.async.db.pool
+
+import java.util.concurrent.{ScheduledFuture, TimeoutException}
+import com.github.mauricio.async.db.util.{ByteBufferUtils, ExecutorServiceUtils}
+import org.specs2.mutable.SpecificationWithJUnit
+import scala.concurrent.duration._
+import scala.concurrent.{Future, Promise}
+
+/**
+ * Tests for TimeoutScheduler
+ */
+class TimeoutSchedulerSpec extends SpecificationWithJUnit {
+
+ val TIMEOUT_DID_NOT_PASS = "timeout did not pass"
+
+
+ "test timeout did not pass" in {
+ val timeoutScheduler = new DummyTimeoutScheduler()
+ val promise = Promise[String]()
+ val scheduledFuture = timeoutScheduler.addTimeout(promise,Duration(1000, MILLISECONDS))
+ Thread.sleep(100);
+ promise.isCompleted === false
+ promise.success(TIMEOUT_DID_NOT_PASS)
+ Thread.sleep(1500)
+ promise.future.value.get.get === TIMEOUT_DID_NOT_PASS
+ scheduledFuture.get.isCancelled === true
+ timeoutScheduler.timeoutCount === 0
+ }
+
+ "test timeout passed" in {
+ val timeoutMillis = 100
+ val promise = Promise[String]()
+ val timeoutScheduler = new DummyTimeoutScheduler()
+ val scheduledFuture = timeoutScheduler.addTimeout(promise,Duration(timeoutMillis, MILLISECONDS))
+ Thread.sleep(1000)
+ promise.isCompleted === true
+ scheduledFuture.get.isCancelled === false
+ promise.trySuccess(TIMEOUT_DID_NOT_PASS)
+ timeoutScheduler.timeoutCount === 1
+ promise.future.value.get.get must throwA[TimeoutException](message = s"Operation is timeouted after it took too long to return \\(${timeoutMillis} milliseconds\\)")
+ }
+
+
+ "test no timeout" in {
+ val timeoutScheduler = new DummyTimeoutScheduler()
+ val promise = Promise[String]()
+ val scheduledFuture = timeoutScheduler.addTimeout(promise,Duration.Inf)
+ Thread.sleep(1000)
+ scheduledFuture === None
+ promise.isCompleted === false
+ promise.success(TIMEOUT_DID_NOT_PASS)
+ promise.future.value.get.get === TIMEOUT_DID_NOT_PASS
+ timeoutScheduler.timeoutCount === 0
+ }
+}
+

mysql-async/src/main/scala/com/github/mauricio/async/db/mysql/MySQLConnection.scala

@@ -25,6 +25,7 @@ import com.github.mauricio.async.db.mysql.exceptions.MySQLException
 import com.github.mauricio.async.db.mysql.message.client._
 import com.github.mauricio.async.db.mysql.message.server._
 import com.github.mauricio.async.db.mysql.util.CharsetMapper
+import com.github.mauricio.async.db.pool.TimeoutScheduler
 import com.github.mauricio.async.db.util.ChannelFutureTransformer.toFuture
 import com.github.mauricio.async.db.util._
 import io.netty.channel.{ChannelHandlerContext, EventLoopGroup}
@@ -46,6 +47,7 @@ class MySQLConnection(
  )
  extends MySQLHandlerDelegate
  with Connection
+ with TimeoutScheduler
 {
 
  import MySQLConnection.log
@@ -56,7 +58,7 @@ class MySQLConnection(
 
  private final val connectionCount = MySQLConnection.Counter.incrementAndGet()
  private final val connectionId = s"[mysql-connection-$connectionCount]"
- private implicit val internalPool = executionContext
+ override implicit val internalPool = executionContext
 
  private final val connectionHandler = new MySQLConnectionHandler(
  configuration,
@@ -188,7 +190,7 @@ class MySQLConnection(
  val promise = Promise[QueryResult]()
  this.setQueryPromise(promise)
  this.connectionHandler.write(new QueryMessage(query))
-
+ addTimeout(promise, configuration.queryTimeout)
  promise.future
  }
 
@@ -224,6 +226,7 @@ class MySQLConnection(
  }
 
  def disconnect: Future[Connection] = this.close
+ override def onTimeout = disconnect
 
  def isConnected: Boolean = this.connectionHandler.isConnected
 
@@ -236,7 +239,7 @@ class MySQLConnection(
  val promise = Promise[QueryResult]()
  this.setQueryPromise(promise)
  this.connectionHandler.sendPreparedStatement(query, values)
-
+ addTimeout(promise,configuration.queryTimeout)
  promise.future
  }
 

mysql-async/src/main/scala/com/github/mauricio/async/db/mysql/pool/MySQLConnectionFactory.scala

@@ -21,9 +21,8 @@ import com.github.mauricio.async.db.pool.ObjectFactory
 import com.github.mauricio.async.db.mysql.MySQLConnection
 import scala.util.Try
 import scala.concurrent.Await
-import scala.concurrent.duration._
 import com.github.mauricio.async.db.util.Log
-import com.github.mauricio.async.db.exceptions.{ConnectionStillRunningQueryException, ConnectionNotConnectedException}
+import com.github.mauricio.async.db.exceptions.{ConnectionTimeoutedException, ConnectionStillRunningQueryException, ConnectionNotConnectedException}
 
 object MySQLConnectionFactory {
  final val log = Log.get[MySQLConnectionFactory]
@@ -90,7 +89,9 @@ class MySQLConnectionFactory( configuration : Configuration ) extends ObjectFact
  */
  def validate(item: MySQLConnection): Try[MySQLConnection] = {
  Try{
-
+ if ( item.isTimeouted ) {
+ throw new ConnectionTimeoutedException(item)
+ }
  if ( !item.isConnected ) {
  throw new ConnectionNotConnectedException(item)
  }

mysql-async/src/test/scala/com/github/mauricio/async/db/mysql/ConnectionHelper.scala

@@ -115,6 +115,19 @@ trait ConnectionHelper {
 
  }
 
+ def withConfigurablePool[T]( configuration : Configuration )( fn : (ConnectionPool[MySQLConnection]) => T ) : T = {
+
+ val factory = new MySQLConnectionFactory(configuration)
+ val pool = new ConnectionPool[MySQLConnection](factory, PoolConfiguration.Default)
+
+ try {
+ fn(pool)
+ } finally {
+ awaitFuture( pool.close )
+ }
+
+ }
+
  def withConnection[T]( fn : (MySQLConnection) => T ) : T =
  withConfigurableConnection(this.defaultConfiguration)(fn)
 

mysql-async/src/test/scala/com/github/mauricio/async/db/mysql/QueryTimeoutSpec.scala

@@ -0,0 +1,80 @@
+/*
+ * Copyright 2013 Maurício Linhares
+ *
+ * Maurício Linhares licenses this file to you under the Apache License,
+ * version 2.0 (the "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at:
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+
+package com.github.mauricio.async.db.mysql
+
+import java.util.concurrent.TimeoutException
+import com.github.mauricio.async.db.Configuration
+import org.specs2.execute.{AsResult, Success, ResultExecution}
+import org.specs2.mutable.Specification
+import scala.concurrent.Await
+import scala.concurrent.duration._
+
+class QueryTimeoutSpec extends Specification with ConnectionHelper {
+ implicit def unitAsResult: AsResult[Unit] = new AsResult[Unit] {
+ def asResult(r: =>Unit) =
+ ResultExecution.execute(r)(_ => Success())
+ }
+ "Simple query with 1 nanosec timeout" in {
+ withConfigurablePool(shortTimeoutConfiguration) {
+ pool => {
+ val connection = Await.result(pool.take, Duration(10,SECONDS))
+ connection.isTimeouted === false
+ connection.isConnected === true
+ val queryResultFuture = connection.sendQuery("select sleep(1)")
+ Await.result(queryResultFuture, Duration(10,SECONDS)) must throwA[TimeoutException]()
+ connection.isTimeouted === true
+ Await.ready(pool.giveBack(connection), Duration(10,SECONDS))
+ pool.availables.count(_ == connection) === 0 // connection removed from pool
+ // we do not know when the connection will be closed.
+ }
+ }
+ }
+
+ "Simple query with 5 sec timeout" in {
+ withConfigurablePool(longTimeoutConfiguration) {
+ pool => {
+ val connection = Await.result(pool.take, Duration(10,SECONDS))
+ connection.isTimeouted === false
+ connection.isConnected === true
+ val queryResultFuture = connection.sendQuery("select sleep(1)")
+ Await.result(queryResultFuture, Duration(10,SECONDS)).rows.get.size === 1
+ connection.isTimeouted === false
+ connection.isConnected === true
+ Await.ready(pool.giveBack(connection), Duration(10,SECONDS))
+ pool.availables.count(_ == connection) === 1 // connection returned to pool
+ }
+ }
+ }
+
+ def shortTimeoutConfiguration = new Configuration(
+ "mysql_async",
+ "localhost",
+ port = 3306,
+ password = Some("root"),
+ database = Some("mysql_async_tests"),
+ queryTimeout = Duration(1,NANOSECONDS)
+ )
+
+ def longTimeoutConfiguration = new Configuration(
+ "mysql_async",
+ "localhost",
+ port = 3306,
+ password = Some("root"),
+ database = Some("mysql_async_tests"),
+ queryTimeout = Duration(5,SECONDS)
+ )
+}

postgresql-async/src/main/scala/com/github/mauricio/async/db/postgresql/PostgreSQLConnection.scala

@@ -20,6 +20,7 @@ import com.github.mauricio.async.db.QueryResult
 import com.github.mauricio.async.db.column.{ColumnEncoderRegistry, ColumnDecoderRegistry}
 import com.github.mauricio.async.db.exceptions.{InsufficientParametersException, ConnectionStillRunningQueryException}
 import com.github.mauricio.async.db.general.MutableResultSet
+import com.github.mauricio.async.db.pool.TimeoutScheduler
 import com.github.mauricio.async.db.postgresql.codec.{PostgreSQLConnectionDelegate, PostgreSQLConnectionHandler}
 import com.github.mauricio.async.db.postgresql.column.{PostgreSQLColumnDecoderRegistry, PostgreSQLColumnEncoderRegistry}
 import com.github.mauricio.async.db.postgresql.exceptions._
@@ -48,7 +49,8 @@ class PostgreSQLConnection
  executionContext : ExecutionContext = ExecutorServiceUtils.CachedExecutionContext
  )
  extends PostgreSQLConnectionDelegate
- with Connection {
+ with Connection
+ with TimeoutScheduler {
 
  import PostgreSQLConnection._
 
@@ -63,7 +65,7 @@ class PostgreSQLConnection
 
  private final val currentCount = Counter.incrementAndGet()
  private final val preparedStatementsCounter = new AtomicInteger()
- private final implicit val internalExecutionContext = executionContext
+ override implicit val internalPool = executionContext
 
  private val parameterStatus = new scala.collection.mutable.HashMap[String, String]()
  private val parsedStatements = new scala.collection.mutable.HashMap[String, PreparedStatementHolder]()
@@ -91,6 +93,7 @@ class PostgreSQLConnection
  }
 
  override def disconnect: Future[Connection] = this.connectionHandler.disconnect.map( c => this )
+ override def onTimeout = disconnect
 
  override def isConnected: Boolean = this.connectionHandler.isConnected
 
@@ -103,7 +106,7 @@ class PostgreSQLConnection
  this.setQueryPromise(promise)
 
  write(new QueryMessage(query))
-
+ addTimeout(promise,configuration.queryTimeout)
  promise.future
  }
 
@@ -130,7 +133,7 @@ class PostgreSQLConnection
  holder.prepared = true
  new PreparedStatementOpeningMessage(holder.statementId, holder.realQuery, values, this.encoderRegistry)
  })
-
+ addTimeout(promise,configuration.queryTimeout)
  promise.future
  }