From ea0a360351dc48c669b995b9c2b077a830daefe3 Mon Sep 17 00:00:00 2001
From: aegbert <aegbert@lucidchart.com>
Date: Wed, 6 Nov 2024 19:20:38 -0700
Subject: [PATCH] Create new connection provider during failovers

---
 .../lucidchart/piezo/ConnectionProvider.scala | 26 +++++++++++++++++--
 1 file changed, 24 insertions(+), 2 deletions(-)

diff --git a/worker/src/main/scala/com/lucidchart/piezo/ConnectionProvider.scala b/worker/src/main/scala/com/lucidchart/piezo/ConnectionProvider.scala
index 6d9be51..1099b77 100644
--- a/worker/src/main/scala/com/lucidchart/piezo/ConnectionProvider.scala
+++ b/worker/src/main/scala/com/lucidchart/piezo/ConnectionProvider.scala
@@ -34,7 +34,8 @@ class ConnectionProvider(props: Properties) {
   @volatile
   private var cachedIpWithExpiration: Option[CachedIpWithExpiration] = None
 
-  private val pool: Pool = new Pool(getIP)
+  @volatile
+  private var pool: Pool = new Pool(getIP)
 
   // Intended to be used only for tests. This mocks an IP failover every time a connection is retreived
   private val causeFailoverEveryConnection = props.getProperty("causeFailoverEveryConnection") == "true"
@@ -111,7 +112,28 @@ class ConnectionProvider(props: Properties) {
       val newIP: String = getIP
       if (hasIpAddressChanged(pool, newIP)) {
         // A failover has occurred, so we evict connections softly. New connectons look up the new IP address
-        pool.connectionProvider.map(_.getDataSource().getHikariPoolMXBean().softEvictConnections())
+        logger.info(s"IP Address has changed for ${jdbcURL}: ${pool.ip} -> ${newIP}. Attempt replacing pool...")
+        val optionalOldPool = synchronized {
+          val oldPool = pool
+          // check if another thread updated the pool
+          if (hasIpAddressChanged(pool, newIP)) {
+            logger.info(s"Replacing pool for ${jdbcURL}...")
+            pool = new Pool(newIP)
+            Some(oldPool)
+          } else {
+            // already up to date
+            logger.info(s"Pool already replaced for ${jdbcURL}")
+            None
+          }
+        }
+
+        // Clean up old pool so we don't leak connections to the old server
+        optionalOldPool.foreach { oldPool =>
+          oldPool.connectionProvider.foreach { provider =>
+            logger.info(s"Closing DB connection pool for ${jdbcURL} for failover (${oldPool.ip} -> ${pool.ip})")
+            provider.shutdown()
+          }
+        }
       }
     }
     pool.connectionProvider.get.getConnection