Merge pull request #63 from ing-bank/fix/redis-connection-pool

Fix/redis connection pool

src/main/scala/com/ing/wbaa/rokku/sts/RokkuStsService.scala

@@ -13,6 +13,7 @@ import com.typesafe.scalalogging.LazyLogging
 import scala.concurrent.duration._
 import scala.concurrent.{ ExecutionContext, Future }
 import scala.util.{ Failure, Success }
+import scala.collection.mutable.ListBuffer
 
 trait RokkuStsService
   extends LazyLogging
@@ -27,6 +28,8 @@ trait RokkuStsService
 
   protected[this] def httpSettings: HttpSettings
 
+  private final val terminationCallbacks: ListBuffer[() => Unit] = ListBuffer()
+
   // The routes we serve
   final val allRoutes: Route =
     toStrictEntity(3.seconds) {
@@ -40,8 +43,14 @@ trait RokkuStsService
 
     Http().newServerAt(httpSettings.httpBind, httpSettings.httpPort).bind(allRoutes)
       .andThen {
-        case Success(binding) => logger.info(s"Sts service started listening: ${binding.localAddress}")
-        case Failure(reason)  => logger.error("Sts service failed to start.", reason)
+        case Success(binding) =>
+          logger.info(s"Sts service started listening: ${binding.localAddress}")
+          sys.addShutdownHook {
+            logger.info("Received termination signal")
+            terminationCallbacks.foreach(c => c())
+            shutdown()
+          }
+        case Failure(reason) => logger.error("Sts service failed to start.", reason)
       }
   }
 
@@ -52,4 +61,8 @@ trait RokkuStsService
         case Failure(reason) => logger.error("Sts service failed to stop.", reason)
       }
   }
+
+  def registerTerminationCallback(f: () => Unit): Unit = {
+    terminationCallbacks += f
+  }
 }

src/main/scala/com/ing/wbaa/rokku/sts/Server.scala

@@ -22,7 +22,19 @@ object Server extends App {
 
     override protected[this] def redisSettings: RedisSettings = RedisSettings(system)
 
+    registerTerminationCallback(() => {
+      logger.info("Closing redis connection pool...")
+      redisPooledConnection.close()
+    })
+
     //Connects to Redis on startup and initializes indeces
-    initializeUserSearchIndex(redisPooledConnection)
+    try {
+      initializeUserSearchIndex(redisPooledConnection)
+    } catch {
+      case ex: RedisSecondaryIndexException =>
+        logger.error(s"Unable to create index $UsersIndex. Error: ${ex.getMessage()}. Exiting...")
+        sys.exit(1)
+    }
+
   }.startup
 }

src/main/scala/com/ing/wbaa/rokku/sts/service/db/Redis.scala

@@ -59,7 +59,9 @@ trait Redis extends LazyLogging {
    */
   protected[this] final def checkDbConnection(): Future[Unit] = {
     Future {
-      val response = new Jedis(redisPooledConnection.getPool().getResource()).ping()
+      val connection = redisPooledConnection.getPool().getResource()
+      val response = new Jedis(connection).ping()
+      connection.close()
       assert(response.toLowerCase().equals("pong"))
     }
   }

src/main/scala/com/ing/wbaa/rokku/sts/service/db/RedisModel.scala

@@ -21,6 +21,8 @@ trait RedisModel extends Encryption {
 
   private val DuplicateIndexExceptionMsg = "Index already exists"
 
+  case class RedisSecondaryIndexException(message: String) extends Exception
+
   object UserKey {
     def encode(username: Username): String = {
       s"$UserKeyPrefix${username.value}"
@@ -89,7 +91,7 @@ trait RedisModel extends Encryption {
           case DuplicateIndexExceptionMsg =>
             logger.info(s"Index ${UsersIndex} already exists. Continuing...")
           case _ =>
-            logger.error(s"Unable to create index $UsersIndex. Error: ${exc.getMessage()}")
+            throw new RedisSecondaryIndexException(s"Unable to create index $UsersIndex. Error: ${exc.getMessage()}")
         }
       case exc: Exception =>
         logger.error(s"Unable to create index $UsersIndex. Error: ${exc.getMessage()}")