Implement per-thread connection caching in BlockingDatabaseWrapper

querulous-core/src/main/scala/com/twitter/querulous/async/BlockingDatabaseWrapper.scala

@@ -45,6 +45,13 @@ extends AsyncDatabase {
       workPoolSize, new DaemonThreadFactory("asyncWorkPool-" + database.hosts.mkString(","))))
   private val openTimeout = database.openTimeout
 
+  // We cache the connection checked out from the underlying database in a thread local so
+  // that each workPool thread can hold on to a connection for its lifetime. This saves expensive
+  // context switches in borrowing/returning connections from the underlying database per request.
+  private val tlConnection = new ThreadLocal[Connection] {
+    override def initialValue() = database.open()
+  }
+
   // Basically all we need to do is offload the real work to workPool. However, there is one
   // complication - enforcement of DB open timeout. If a connection is not available, most
   // likely neither is a thread to do the work, so requests would queue up in the future pool.
@@ -57,11 +64,20 @@ extends AsyncDatabase {
     val future = workPool {
       val isRunnable = startCoordinator.compareAndSet(true, false)
       if (isRunnable) {
-        val connection = database.open()
+        val connection = tlConnection.get()
         try {
           f(connection)
-        } finally {
-          database.close(connection)
+        } catch {
+          case e => {
+            // An exception occurred. To be safe, we return our cached connection back to the pool. This
+            // protects us in case either the connection has been killed or our thread is going to be
+            // terminated with an unhandled exception. If neither is the case (e.g. this was a benign
+            // exception like a SQL constraint violation), it still doesn't hurt much to return/re-borrow
+            // the connection from the underlying database, given that this should be rare.
+            database.close(connection)
+            tlConnection.remove()
+            throw e
+          }
         }
       } else {
         throw new CancellationException

querulous-core/src/test/scala/com/twitter/querulous/unit/BlockingDatabaseWrapperSpec.scala

@@ -32,13 +32,21 @@ class BlockingDatabaseWrapperSpec extends Specification {
       def reset() { totalOpens.set(0); openConns.set(0) }
     }
 
-    val wrapper = new BlockingDatabaseWrapper(1, database)
+    val numThreads = 2
+    val wrapper = new BlockingDatabaseWrapper(numThreads, database)
 
     doBefore { database.reset() }
 
     "withConnection should follow connection lifecycle" in {
       wrapper withConnection { _ => "Done" } apply()
 
+      database.totalOpens.get mustEqual 1
+      database.openConns.get  mustEqual 1
+    }
+
+    "withConnection should return connection on exception" in {
+      wrapper withConnection { _ => throw new Exception } handle { case _ => "Done with Exception" } apply()
+
       database.totalOpens.get mustEqual 1
       database.openConns.get  mustEqual 0
     }
@@ -51,7 +59,7 @@ class BlockingDatabaseWrapperSpec extends Specification {
 
       result mustBe "Done"
       database.totalOpens.get mustEqual 1
-      database.openConns.get  mustEqual 0
+      database.openConns.get  mustEqual 1
     }
 
     "withConnection should follow lifecycle regardless of cancellation" in {
@@ -77,10 +85,10 @@ class BlockingDatabaseWrapperSpec extends Specification {
       println("Ran block: "+ hitBlock.get)
       println("Cancelled: "+ (100000 - completed.size))
       println("Completed: "+ completed.size)
-      println("Leaked:    "+ database.openConns.get)
+      println("Cached:    "+ database.openConns.get)
 
-      database.totalOpens.get mustEqual completed.size
-      database.openConns.get mustEqual 0
+      database.totalOpens.get mustEqual numThreads
+      database.openConns.get mustEqual numThreads
     }
   }
 }