allow hikari password changes at runtime

atlasdb-dbkvs-hikari/src/main/java/com/palantir/nexus/db/pool/HikariCPConnectionManager.java

@@ -23,6 +23,7 @@
 import com.palantir.nexus.db.DBType;
 import com.palantir.nexus.db.pool.config.ConnectionConfig;
 import com.palantir.nexus.db.sql.ExceptionCheck;
+import com.zaxxer.hikari.HikariConfig;
 import com.zaxxer.hikari.HikariDataSource;
 import com.zaxxer.hikari.HikariPoolMXBean;
 import com.zaxxer.hikari.pool.HikariPool.PoolInitializationException;
@@ -52,7 +53,8 @@
 public class HikariCPConnectionManager extends BaseConnectionManager {
     private static final Logger log = LoggerFactory.getLogger(HikariCPConnectionManager.class);
 
-    private ConnectionConfig connConfig;
+    private final ConnectionConfig connConfig;
+    private final HikariConfig hikariConfig;
 
     private enum StateType {
         // Base state at construction.  Nothing is set.
@@ -83,6 +85,7 @@ private static class State {
 
     public HikariCPConnectionManager(ConnectionConfig connConfig) {
         this.connConfig = Preconditions.checkNotNull(connConfig, "ConnectionConfig must not be null");
+        this.hikariConfig = connConfig.getHikariConfig();
     }
 
     @Override
@@ -293,17 +296,15 @@ private HikariDataSource getDataSourcePool() {
 
         try {
             try {
-                dataSourcePool = new HikariDataSource(connConfig.getHikariConfig());
+                dataSourcePool = new HikariDataSource(hikariConfig);
             } catch (IllegalArgumentException e) {
                 // allow multiple pools on same JVM (they need unique names / endpoints)
                 if (e.getMessage().contains("A metric named")) {
                     String poolName = connConfig.getConnectionPoolName();
 
-                    connConfig
-                            .getHikariConfig()
-                            .setPoolName(
-                                    poolName + "-" + ThreadLocalRandom.current().nextInt());
-                    dataSourcePool = new HikariDataSource(connConfig.getHikariConfig());
+                    hikariConfig.setPoolName(
+                            poolName + "-" + ThreadLocalRandom.current().nextInt());
+                    dataSourcePool = new HikariDataSource(hikariConfig);
                 } else {
                     throw e;
                 }
@@ -364,6 +365,50 @@ public DBType getDbType() {
         return connConfig.getDbType();
     }
 
+    @Override
+    public void setPassword(String newPassword) {
+        Preconditions.checkNotNull(newPassword, "password cannot be null");
+        State localState = state;
+        if (localState.type == StateType.ZERO) {
+            synchronized (this) {
+                // need to check again in case init just happened
+                localState = state;
+                if (localState.type == StateType.ZERO) {
+                    // the pool is not initialized, need to set password on the config object
+                    internalSetPassword(newPassword, hikariConfig);
+                    return;
+                }
+            }
+        }
+        if (localState.type == StateType.NORMAL) {
+            // pool is initialized, need to set password directly on the pool
+            internalSetPassword(newPassword, localState.dataSourcePool);
+        } else if (localState.type == StateType.CLOSED) {
+            throw new SafeRuntimeException("Hikari pool is closed", localState.closeTrace);
+        }
+    }
+
+    /**
+     * Sets the password (and username) on the HikariConfig (or HikariDataSource pool) if it is necessary.
+     *
+     * <p>Note that when the password matches static config, performance is better when setting the username and
+     * password to null so that hikari uses the original JDBC parameters and avoids making new copies for every
+     * new connection. See {@link com.zaxxer.hikari.pool.PoolBase#newConnection}. Also see
+     * {@link com.palantir.nexus.db.pool.config.OracleConnectionConfig#getHikariProperties} for the JDBC properties.
+     *
+     * <p>Hikari checks the username field for being null to determine if it needs to copy the JDBC parameters, so
+     * the username *must* be set for it to respect the password change.
+     */
+    private void internalSetPassword(String newPassword, HikariConfig configOrPool) {
+        if (newPassword.equals(connConfig.getDbPassword().unmasked())) {
+            configOrPool.setUsername(null);
+            configOrPool.setPassword(null);
+        } else {
+            configOrPool.setPassword(newPassword);
+            configOrPool.setUsername(connConfig.getDbLogin());
+        }
+    }
+
     private final class ConnectionAcquisitionProfiler {
         private final Stopwatch globalStopwatch;
         private final Stopwatch trialStopwatch;

atlasdb-dbkvs-tests/src/test/java/com/palantir/atlasdb/keyvalue/dbkvs/DbkvsPostgresTestSuite.java

@@ -23,9 +23,10 @@
 import com.palantir.docker.compose.connection.Container;
 import com.palantir.docker.compose.connection.DockerPort;
 import com.palantir.docker.compose.logging.LogDirectory;
-import com.palantir.nexus.db.pool.config.ConnectionConfig;
+import com.palantir.nexus.db.pool.HikariCpConnectionManagerTest;
 import com.palantir.nexus.db.pool.config.ImmutableMaskedValue;
 import com.palantir.nexus.db.pool.config.ImmutablePostgresConnectionConfig;
+import com.palantir.nexus.db.pool.config.PostgresConnectionConfig;
 import java.net.InetSocketAddress;
 import java.time.Duration;
 import java.util.concurrent.Callable;
@@ -49,7 +50,8 @@
     DbKvsPostgresGetCandidateCellsForSweepingTest.class,
     DbKvsSweepProgressStoreIntegrationTest.class,
     DbKvsPostgresInvalidationRunnerTest.class,
-    DbTimestampStoreInvalidatorCreationTest.class
+    DbTimestampStoreInvalidatorCreationTest.class,
+    HikariCpConnectionManagerTest.class,
 })
 public final class DbkvsPostgresTestSuite {
     private static final int POSTGRES_PORT_NUMBER = 5432;
@@ -74,21 +76,21 @@ public static void waitUntilDbkvsIsUp() throws InterruptedException {
                 .until(canCreateKeyValueService());
     }
 
-    public static DbKeyValueServiceConfig getKvsConfig() {
+    public static PostgresConnectionConfig getConnectionConfig() {
         DockerPort port = docker.containers().container("postgres").port(POSTGRES_PORT_NUMBER);
-
         InetSocketAddress postgresAddress = new InetSocketAddress(port.getIp(), port.getExternalPort());
-
-        ConnectionConfig connectionConfig = ImmutablePostgresConnectionConfig.builder()
+        return ImmutablePostgresConnectionConfig.builder()
                 .dbName("atlas")
                 .dbLogin("palantir")
                 .dbPassword(ImmutableMaskedValue.of("palantir"))
                 .host(postgresAddress.getHostName())
                 .port(postgresAddress.getPort())
                 .build();
+    }
 
+    public static DbKeyValueServiceConfig getKvsConfig() {
         return ImmutableDbKeyValueServiceConfig.builder()
-                .connection(connectionConfig)
+                .connection(getConnectionConfig())
                 .ddl(ImmutablePostgresDdlConfig.builder()
                         .compactInterval(HumanReadableDuration.days(2))
                         .build())

atlasdb-dbkvs-tests/src/test/java/com/palantir/nexus/db/pool/HikariCpConnectionManagerTest.java

@@ -19,39 +19,28 @@
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.fail;
 
-import com.palantir.docker.compose.DockerComposeRule;
-import com.palantir.docker.compose.configuration.ShutdownStrategy;
-import com.palantir.docker.compose.connection.Container;
-import com.palantir.docker.compose.connection.DockerPort;
-import com.palantir.docker.compose.logging.LogDirectory;
+import com.google.common.base.Throwables;
+import com.palantir.atlasdb.keyvalue.dbkvs.DbkvsPostgresTestSuite;
+import com.palantir.logsafe.exceptions.SafeRuntimeException;
 import com.palantir.nexus.db.pool.config.ConnectionConfig;
 import com.palantir.nexus.db.pool.config.ImmutableMaskedValue;
 import com.palantir.nexus.db.pool.config.ImmutablePostgresConnectionConfig;
-import java.net.InetSocketAddress;
+import com.palantir.nexus.db.pool.config.PostgresConnectionConfig;
 import java.sql.Connection;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.SQLTransientConnectionException;
 import java.sql.Statement;
+import java.time.Duration;
+import org.awaitility.Awaitility;
 import org.junit.After;
 import org.junit.Before;
-import org.junit.ClassRule;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
 
 public class HikariCpConnectionManagerTest {
 
-    private static final int POSTGRES_PORT_NUMBER = 5432;
-
-    @ClassRule
-    public static final DockerComposeRule docker = DockerComposeRule.builder()
-            .file("src/test/resources/docker-compose.yml")
-            .waitingForService("postgres", Container::areAllPortsOpen)
-            .saveLogsTo(LogDirectory.circleAwareLogDirectory(HikariCpConnectionManagerTest.class))
-            .shutdownStrategy(ShutdownStrategy.AGGRESSIVE_WITH_NETWORK_CLEANUP)
-            .build();
-
     private ConnectionManager manager;
 
     @Before
@@ -141,6 +130,103 @@ public void testCloseIdempotent() throws SQLException {
         manager.close(); // shouldn't throw
     }
 
+    @Test
+    public void testRuntimePasswordChange() throws SQLException {
+        // create a new user to avoid messing up the main user for other tests
+        String testUsername = "testpasswordchange";
+        String password1 = "password1";
+        String password2 = "password2";
+        String password3 = "password3";
+
+        // initial config intentionally has the wrong password for the test user
+        PostgresConnectionConfig testConfig = createConnectionConfig(testUsername, password3, 0, 3);
+
+        try (Connection conn = manager.getConnection()) {
+            try (Statement statement = conn.createStatement()) {
+                statement.execute(String.format("CREATE USER %s WITH PASSWORD '%s'", testUsername, password1));
+            }
+            try (Statement statement = conn.createStatement()) {
+                statement.execute(String.format(
+                        "GRANT ALL PRIVILEGES ON DATABASE %s TO %s", testConfig.getDbName(), testUsername));
+            }
+        }
+
+        ConnectionManager testManager = new HikariCPConnectionManager(testConfig);
+        // fixing the password before init should work
+        testManager.setPassword(password1);
+
+        try (Connection conn = testManager.getConnection()) {
+            try (Statement statement = conn.createStatement()) {
+                statement.execute(String.format("ALTER USER %s WITH PASSWORD '%s'", testUsername, password2));
+            }
+            // existing connection should still work
+            checkConnection(conn);
+
+            // trying to get a new connection should fail (times out with wrong password)
+            assertPasswordWrong(testManager);
+
+            // fix the password on the pool
+            testManager.setPassword(password2);
+            // original connection should still work
+            checkConnection(conn);
+
+            // new connection should also work
+            try (Connection conn2 = testManager.getConnection()) {
+                checkConnection(conn2);
+            }
+
+            // changing the pool password again and one connection should work (conn still in the pool)
+            testManager.setPassword(password3);
+            try (Connection conn2 = testManager.getConnection()) {
+                checkConnection(conn2);
+
+                // getting one more connection should see the password error
+                assertPasswordWrong(testManager);
+
+                // use existing conn to fix the password
+                try (Statement statement = conn2.createStatement()) {
+                    statement.execute(String.format("ALTER USER %s WITH PASSWORD '%s'", testUsername, password3));
+                }
+
+                // now a new connection should work
+                try (Connection conn3 = testManager.getConnection()) {
+                    checkConnection(conn3);
+                }
+            }
+        } finally {
+            testManager.close();
+        }
+    }
+
+    private static void assertPasswordWrong(ConnectionManager testManager) {
+        // This is needed because it appears that sometimes postgres password changes do not take effect immediately.
+        // If the password change happens too late, we end up with a connection in the pool that we did not expect.
+        // In that case we need to wait the configured time (1 second) for hikari to remove the idle connection from
+        // the pool before we can detect that the password is wrong. Note that I cannot reproduce this case locally,
+        // but it appears to almost always happen on circle.
+        Awaitility.await("assertPasswordWrong")
+                .atMost(Duration.ofSeconds(10))
+                .pollInterval(Duration.ofSeconds(2))
+                .pollDelay(Duration.ofMillis(100))
+                .until(() -> isPasswordWrong(testManager));
+    }
+
+    private static boolean isPasswordWrong(ConnectionManager testManager) {
+        try (Connection ignored = testManager.getConnection()) {
+            return false;
+        } catch (SQLException e) {
+            // if "password authentication failed" is in the cause chain, the password is wrong
+            // otherwise this is an unexpected exception
+            for (Throwable t : Throwables.getCausalChain(e)) {
+                String message = t.getMessage();
+                if (message != null && message.contains("password authentication failed")) {
+                    return true;
+                }
+            }
+            throw new SafeRuntimeException("unexpected exception checking password on ConnectionManager", e);
+        }
+    }
+
     private static void checkConnection(Connection conn) throws SQLException {
         try (Statement statement = conn.createStatement()) {
             try (ResultSet result = statement.executeQuery("SELECT 123")) {
@@ -152,16 +238,23 @@ private static void checkConnection(Connection conn) throws SQLException {
     }
 
     private static ConnectionConfig createConnectionConfig(int maxConnections) {
-        DockerPort port = docker.containers().container("postgres").port(POSTGRES_PORT_NUMBER);
-        InetSocketAddress postgresAddress = new InetSocketAddress(port.getIp(), port.getExternalPort());
+        return createConnectionConfig("palantir", "palantir", maxConnections, maxConnections);
+    }
+
+    private static PostgresConnectionConfig createConnectionConfig(
+            String username, String password, int minConnections, int maxConnections) {
+        PostgresConnectionConfig suiteConfig = DbkvsPostgresTestSuite.getConnectionConfig();
         return ImmutablePostgresConnectionConfig.builder()
-                .dbName("atlas")
-                .dbLogin("palantir")
-                .dbPassword(ImmutableMaskedValue.of("palantir"))
-                .host(postgresAddress.getHostName())
-                .port(postgresAddress.getPort())
+                .dbName(suiteConfig.getDbName())
+                .host(suiteConfig.getHost())
+                .port(suiteConfig.getPort())
+                .dbLogin(username)
+                .dbPassword(ImmutableMaskedValue.of(password))
+                .minConnections(minConnections)
                 .maxConnections(maxConnections)
                 .checkoutTimeout(2000)
+                // if this is set too high, testRuntimePasswordChange will be unreliable (and will take longer)
+                .maxConnectionAge(1)
                 .build();
     }
 }

changelog/@unreleased/pr-5308.v2.yml

@@ -0,0 +1,6 @@
+type: improvement
+improvement:
+  description: |-
+    The Hikari ConnectionManager now has a setPassword method which allows changing the password at runtime without needing to shutdown and restart services. A future change will add standard runtime configs for this (for now this is only usable by internal products that have their own custom runtime config).
+  links:
+  - https://github.com/palantir/atlasdb/pull/5308

commons-db/src/main/java/com/palantir/nexus/db/pool/ConnectionManager.java

@@ -51,4 +51,16 @@ public interface ConnectionManager {
     void initUnchecked();
 
     DBType getDbType();
+
+    /**
+     * Change the password which will be used for new connections. Existing connections are not modified in any way.
+     *
+     * <p>Note that if the password is changed at runtime, there is a slight negative impact on pool performance
+     * because Hikari will copy the JDBC parameters every time a new connection is made.
+     *
+     * <p>Concurrent calls to this method should be avoided since the behavior for concurrent calls is not
+     * deterministic. It is recommended to only call this method from a single runtime config subscription so this is
+     * called when the runtime config changes.
+     */
+    void setPassword(String newPassword);
 }

commons-db/src/main/java/com/palantir/nexus/db/pool/ResourceTypes.java

@@ -59,6 +59,11 @@ public void init() throws SQLException {
                         public DBType getDbType() {
                             return delegate.getDbType();
                         }
+
+                        @Override
+                        public void setPassword(String newPassword) {
+                            delegate.setPassword(newPassword);
+                        }
                     };
                 }