Merge pull request mauricio#88 from fwbrasil/partitioned-pool

partitioned db pool

Author: mauricio

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

@@ -0,0 +1,65 @@
+package com.github.mauricio.async.db.pool
+
+import scala.concurrent.Future
+import com.github.mauricio.async.db.util.ExecutorServiceUtils
+import scala.concurrent.Promise
+import java.util.concurrent.ConcurrentHashMap
+import scala.util.Success
+import scala.util.Failure
+
+class PartitionedAsyncObjectPool[T](
+ factory: ObjectFactory[T],
+ configuration: PoolConfiguration,
+ numberOfPartitions: Int)
+ extends AsyncObjectPool[T] {
+
+ import ExecutorServiceUtils.CachedExecutionContext
+
+ private val pools =
+ (0 until numberOfPartitions)
+ .map(_ -> new SingleThreadedAsyncObjectPool(factory, partitionConfig))
+ .toMap
+
+ private val checkouts = new ConcurrentHashMap[T, SingleThreadedAsyncObjectPool[T]]
+
+ def take: Future[T] = {
+ val pool = currentPool
+ pool.take.andThen {
+ case Success(conn) =>
+ checkouts.put(conn, pool)
+ case Failure(_) =>
+ }
+ }
+
+ def giveBack(item: T) =
+ checkouts
+ .remove(item)
+ .giveBack(item)
+ .map(_ => this)
+
+ def close =
+ Future.sequence(pools.values.map(_.close)).map {
+ _ => this
+ }
+
+ def availables: Traversable[T] = pools.values.map(_.availables).flatten
+
+ def inUse: Traversable[T] = pools.values.map(_.inUse).flatten
+
+ def queued: Traversable[Promise[T]] = pools.values.map(_.queued).flatten
+
+ protected def isClosed =
+ pools.values.forall(_.isClosed)
+
+ private def currentPool =
+ pools(currentThreadAffinity)
+
+ private def currentThreadAffinity =
+ (Thread.currentThread.getId % numberOfPartitions).toInt
+
+ private def partitionConfig =
+ configuration.copy(
+ maxObjects = configuration.maxObjects / numberOfPartitions,
+ maxQueueSize = configuration.maxQueueSize / numberOfPartitions
+ )
+}

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

@@ -0,0 +1,33 @@
+package com.github.mauricio.async.db.pool;
+
+import com.github.mauricio.async.db.util.ExecutorServiceUtils
+import com.github.mauricio.async.db.{ QueryResult, Connection }
+import scala.concurrent.{ ExecutionContext, Future }
+
+class PartitionedConnectionPool[T <: Connection](
+ factory: ObjectFactory[T],
+ configuration: PoolConfiguration,
+ numberOfPartitions: Int,
+ executionContext: ExecutionContext = ExecutorServiceUtils.CachedExecutionContext)
+ extends PartitionedAsyncObjectPool[T](factory, configuration, numberOfPartitions)
+ with Connection {
+
+ def disconnect: Future[Connection] = if (this.isConnected) {
+ this.close.map(item => this)(executionContext)
+ } else {
+ Future.successful(this)
+ }
+
+ def connect: Future[Connection] = Future.successful(this)
+
+ def isConnected: Boolean = !this.isClosed
+
+ def sendQuery(query: String): Future[QueryResult] =
+ this.use(_.sendQuery(query))(executionContext)
+
+ def sendPreparedStatement(query: String, values: Seq[Any] = List()): Future[QueryResult] =
+ this.use(_.sendPreparedStatement(query, values))(executionContext)
+
+ override def inTransaction[A](f: Connection => Future[A])(implicit context: ExecutionContext = executionContext): Future[A] =
+ this.use(_.inTransaction[A](f)(context))(executionContext)
+}

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

@@ -0,0 +1,288 @@
+package com.github.mauricio.async.db.pool
+
+import org.specs2.mutable.Specification
+import scala.util.Try
+import scala.concurrent.Await
+import scala.concurrent.duration._
+import scala.concurrent.Future
+import org.specs2.mutable.SpecificationWithJUnit
+import language.reflectiveCalls
+import com.github.mauricio.async.db.util.ExecutorServiceUtils
+import scala.concurrent.ExecutionContext
+import java.util.concurrent.Executors
+
+class PartitionedAsyncObjectPoolSpec extends SpecificationWithJUnit {
+ isolated
+ sequential
+
+ val config =
+ PoolConfiguration(100, Long.MaxValue, 100, Int.MaxValue)
+
+ val factory = new ObjectFactory[Int] {
+ var reject = Set[Int]()
+ var failCreate = false
+ private var current = 0
+ def create =
+ if (failCreate)
+ throw new IllegalStateException
+ else {
+ current += 1
+ current
+ }
+ def destroy(item: Int) = {}
+ def validate(item: Int) =
+ Try {
+ if (reject.contains(item))
+ throw new IllegalStateException
+ else item
+ }
+ }
+
+ val pool = new PartitionedAsyncObjectPool(factory, config, 2)
+ def maxObjects = config.maxObjects / 2
+ def maxIdle = config.maxIdle / 2
+ def maxQueueSize = config.maxQueueSize / 2
+
+ "pool contents" >> {
+
+ "before exceed maxObjects" >> {
+
+ "take one element" in {
+ takeAndWait(1)
+
+ pool.inUse.size mustEqual 1
+ pool.queued.size mustEqual 0
+ pool.availables.size mustEqual 0
+ }
+
+ "take one element and return it invalid" in {
+ takeAndWait(1)
+ factory.reject += 1
+
+ await(pool.giveBack(1)) must throwA[IllegalStateException]
+
+ pool.inUse.size mustEqual 0
+ pool.queued.size mustEqual 0
+ pool.availables.size mustEqual 0
+ }
+
+ "take one failed element" in {
+ factory.failCreate = true
+ takeAndWait(1) must throwA[IllegalStateException]
+
+ pool.inUse.size mustEqual 0
+ pool.queued.size mustEqual 0
+ pool.availables.size mustEqual 0
+ }
+
+ "take maxObjects" in {
+ takeAndWait(maxObjects)
+
+ pool.inUse.size mustEqual maxObjects
+ pool.queued.size mustEqual 0
+ pool.availables.size mustEqual 0
+ }
+
+ "take maxObjects - 1 and take one failed" in {
+ takeAndWait(maxObjects - 1)
+
+ factory.failCreate = true
+ takeAndWait(1) must throwA[IllegalStateException]
+
+ pool.inUse.size mustEqual maxObjects - 1
+ pool.queued.size mustEqual 0
+ pool.availables.size mustEqual 0
+ }
+
+ "take maxObjects and receive one back" in {
+ takeAndWait(maxObjects)
+ await(pool.giveBack(1))
+
+ pool.inUse.size mustEqual maxObjects - 1
+ pool.queued.size mustEqual 0
+ pool.availables.size mustEqual 1
+ }
+
+ "take maxObjects and receive one invalid back" in {
+ takeAndWait(maxObjects)
+ factory.reject += 1
+ await(pool.giveBack(1)) must throwA[IllegalStateException]
+
+ pool.inUse.size mustEqual maxObjects - 1
+ pool.queued.size mustEqual 0
+ pool.availables.size mustEqual 0
+ }
+ }
+
+ "after exceed maxObjects" >> {
+
+ takeAndWait(maxObjects)
+
+ "before exceed maxQueueSize" >> {
+
+ "one take queued" in {
+ pool.take
+
+ pool.inUse.size mustEqual maxObjects
+ pool.queued.size mustEqual 1
+ pool.availables.size mustEqual 0
+ }
+
+ "one take queued and receive one item back" in {
+ val taking = pool.take
+
+ await(pool.giveBack(1))
+
+ await(taking) mustEqual 1
+ pool.inUse.size mustEqual maxObjects
+ pool.queued.size mustEqual 0
+ pool.availables.size mustEqual 0
+ }
+
+ "one take queued and receive one invalid item back" in {
+ val taking = pool.take
+ factory.reject += 1
+ await(pool.giveBack(1)) must throwA[IllegalStateException]
+
+ pool.inUse.size mustEqual maxObjects - 1
+ pool.queued.size mustEqual 1
+ pool.availables.size mustEqual 0
+ }
+
+ "maxQueueSize takes queued" in {
+ for (_ <- 0 until maxQueueSize)
+ pool.take
+
+ pool.inUse.size mustEqual maxObjects
+ pool.queued.size mustEqual maxQueueSize
+ pool.availables.size mustEqual 0
+ }
+
+ "maxQueueSize takes queued and receive one back" in {
+ val taking = pool.take
+ for (_ <- 0 until maxQueueSize - 1)
+ pool.take
+
+ await(pool.giveBack(10))
+
+ await(taking) mustEqual 10
+ pool.inUse.size mustEqual maxObjects
+ pool.queued.size mustEqual maxQueueSize - 1
+ pool.availables.size mustEqual 0
+ }
+
+ "maxQueueSize takes queued and receive one invalid back" in {
+ for (_ <- 0 until maxQueueSize)
+ pool.take
+
+ factory.reject += 11
+ await(pool.giveBack(11)) must throwA[IllegalStateException]
+
+ pool.inUse.size mustEqual maxObjects - 1
+ pool.queued.size mustEqual maxQueueSize
+ pool.availables.size mustEqual 0
+ }
+ }
+
+ "after exceed maxQueueSize" >> {
+
+ for (_ <- 0 until maxQueueSize)
+ pool.take
+
+ "start to reject takes" in {
+ await(pool.take) must throwA[PoolExhaustedException]
+
+ pool.inUse.size mustEqual maxObjects
+ pool.queued.size mustEqual maxQueueSize
+ pool.availables.size mustEqual 0
+ }
+
+ "receive an object back" in {
+ await(pool.giveBack(1))
+
+ pool.inUse.size mustEqual maxObjects
+ pool.queued.size mustEqual maxQueueSize - 1
+ pool.availables.size mustEqual 0
+ }
+
+ "receive an invalid object back" in {
+ factory.reject += 1
+ await(pool.giveBack(1)) must throwA[IllegalStateException]
+
+ pool.inUse.size mustEqual maxObjects - 1
+ pool.queued.size mustEqual maxQueueSize
+ pool.availables.size mustEqual 0
+ }
+
+ "receive maxQueueSize objects back" in {
+ for (i <- 1 to maxQueueSize)
+ await(pool.giveBack(i))
+
+ pool.inUse.size mustEqual maxObjects
+ pool.queued.size mustEqual 0
+ pool.availables.size mustEqual 0
+ }
+
+ "receive maxQueueSize invalid objects back" in {
+ for (i <- 1 to maxQueueSize) {
+ factory.reject += i
+ await(pool.giveBack(i)) must throwA[IllegalStateException]
+ }
+
+ pool.inUse.size mustEqual maxObjects - maxQueueSize
+ pool.queued.size mustEqual maxQueueSize
+ pool.availables.size mustEqual 0
+ }
+
+ "receive maxQueueSize + 1 object back" in {
+ for (i <- 1 to maxQueueSize)
+ await(pool.giveBack(i))
+
+ await(pool.giveBack(1))
+ pool.inUse.size mustEqual maxObjects - 1
+ pool.queued.size mustEqual 0
+ pool.availables.size mustEqual 1
+ }
+
+ "receive maxQueueSize + 1 invalid object back" in {
+ for (i <- 1 to maxQueueSize)
+ await(pool.giveBack(i))
+
+ factory.reject += 1
+ await(pool.giveBack(1)) must throwA[IllegalStateException]
+ pool.inUse.size mustEqual maxObjects - 1
+ pool.queued.size mustEqual 0
+ pool.availables.size mustEqual 0
+ }
+ }
+ }
+ }
+
+ "gives back the connection to the original pool" in {
+ val executor = Executors.newFixedThreadPool(20)
+ implicit val context = ExecutionContext.fromExecutor(executor)
+
+ val takes =
+ for (_ <- 0 until 30) yield {
+ Future().flatMap(_ => pool.take)
+ }
+ val takesAndReturns =
+ Future.sequence(takes).flatMap { items =>
+ Future.sequence(items.map(pool.giveBack))
+ }
+
+ await(takesAndReturns)
+
+ executor.shutdown
+ pool.inUse.size mustEqual 0
+ pool.queued.size mustEqual 0
+ pool.availables.size mustEqual 30
+ }
+
+ private def takeAndWait(objects: Int) =
+ for (_ <- 0 until objects)
+ await(pool.take)
+
+ private def await[T](future: Future[T]) =
+ Await.result(future, Duration.Inf)
+}