Skip to content

Commit

Permalink
conn-man: add test for connection pooling
Browse files Browse the repository at this point in the history
  • Loading branch information
@akphi
akphi committed Oct 11, 2023
1 parent 178f82e commit 7d2febd
Show file tree
Hide file tree
Showing 7 changed files with 218 additions and 95 deletions.
13 changes: 7 additions & 6 deletions legend-engine-config/legend-engine-connection-integration-tests/pom.xml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,8 @@
limitations under the License.
-->

<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
<parent>
<groupId>org.finos.legend.engine</groupId>
<artifactId>legend-engine-config</artifactId>
Expand Down Expand Up @@ -74,11 +75,6 @@
</dependency>
<!-- RUNTIME -->

<dependency>
<groupId>com.zaxxer</groupId>
<artifactId>HikariCP</artifactId>
</dependency>

<!-- TEST -->
<dependency>
<groupId>org.junit.jupiter</groupId>
Expand All @@ -90,6 +86,11 @@
<artifactId>junit-jupiter-engine</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>net.bytebuddy</groupId>
<artifactId>byte-buddy</artifactId>
<scope>test</scope>
</dependency>
<!-- TEST -->
</dependencies>
</project>
197 changes: 123 additions & 74 deletions ...ests/src/test/java/org/finos/legend/engine/connection/test/TestJDBCConnectionManager.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,9 +14,7 @@

package org.finos.legend.engine.connection.test;

import com.zaxxer.hikari.HikariConfig;
import com.zaxxer.hikari.HikariPoolMXBean;
import org.eclipse.collections.api.block.function.Function;
import net.bytebuddy.asm.Advice;
import org.finos.legend.authentication.vault.impl.PropertiesFileCredentialVault;
import org.finos.legend.connection.AuthenticationMechanismConfiguration;
import org.finos.legend.connection.Authenticator;
Expand Down Expand Up @@ -44,13 +42,6 @@
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

import javax.management.JMX;
import javax.management.MBeanServer;
import javax.management.MalformedObjectNameException;
import javax.management.ObjectName;
import java.lang.management.ManagementFactory;
import java.lang.reflect.ParameterizedType;
import java.lang.reflect.Type;
import java.sql.Connection;
import java.util.Properties;

Expand Down Expand Up @@ -107,6 +98,8 @@ public void setup()
public void cleanUp()
{
postgresContainer.stop();

JDBCConnectionManager.getInstance().flushPool();
}

@Test
Expand All @@ -117,14 +110,9 @@ public void testBasicConnectionPooling() throws Exception
new UserPasswordCredentialBuilder()
)
.withConnectionBuilders(
new InstrumentedStaticJDBCConnectionBuilder.WithPlaintextUsernamePassword((HikariConfig config) -> {
config.setKeepaliveTime(10000);
config.setMaximumPoolSize(10);
return null;
})
new StaticJDBCConnectionBuilder.WithPlaintextUsernamePassword()
)
.build();

this.storeInstanceProvider.injectStoreInstance(this.storeInstance);
Identity identity = identityFactory.createIdentity(
new IdentitySpecification.Builder()
Expand All @@ -136,91 +124,152 @@ public void testBasicConnectionPooling() throws Exception
postgresContainer.getUser(),
new PropertiesFileSecret("passwordRef")
);

Authenticator authenticator = this.connectionFactory.getAuthenticator(identity, TEST_STORE_INSTANCE_NAME, authenticationConfiguration);

JDBCConnectionManager connectionManager = JDBCConnectionManager.getInstance();
Assertions.assertEquals(0, connectionManager.getPoolSize());

// 1. Get a connection, this should initialize the pool as well as create a new connection in the empty pool
// this connection should be active
Connection connection0 = this.connectionFactory.getConnection(identity, authenticator);

HikariPoolMXBean poolProxy = getPoolProxy(identity, connectionSpecification, authenticationConfiguration);
Assertions.assertEquals(1, poolProxy.getTotalConnections());
Assertions.assertEquals(1, poolProxy.getActiveConnections());
Assertions.assertEquals(0, poolProxy.getIdleConnections());
String poolName = JDBCConnectionManager.getPoolName(identity, connectionSpecification, authenticationConfiguration);
JDBCConnectionManager.ConnectionPool connectionPool = connectionManager.getPool(poolName);

// 2. Close the connection, verify that the pool keeps this connection around in idle state
Connection underlyingConnection0 = connection0.unwrap(Connection.class);
connection0.close();

Assertions.assertEquals(1, poolProxy.getTotalConnections());
Assertions.assertEquals(0, poolProxy.getActiveConnections());
Assertions.assertEquals(1, poolProxy.getIdleConnections());
Assertions.assertEquals(1, connectionPool.getTotalConnections());
Assertions.assertEquals(0, connectionPool.getActiveConnections());
Assertions.assertEquals(1, connectionPool.getIdleConnections());

// 3. Get a new connection, the pool should return the idle connection and create no new connection
Connection connection1 = this.connectionFactory.getConnection(identity, authenticator);
// Connection connection2 = this.connectionFactory.getConnection(identity, authenticator);

// connection1.isValid()
Assertions.assertEquals(1, poolProxy.getTotalConnections());
Assertions.assertEquals(1, poolProxy.getActiveConnections());
Assertions.assertEquals(0, poolProxy.getIdleConnections());
Assertions.assertEquals(underlyingConnection0, connection1.unwrap(Connection.class));
Assertions.assertEquals(1, connectionPool.getTotalConnections());
Assertions.assertEquals(1, connectionPool.getActiveConnections());
Assertions.assertEquals(0, connectionPool.getIdleConnections());

// 4. Get another connection while the first one is still alive and used, a new connection
// will be created in the pool
this.connectionFactory.getConnection(identity, authenticator);

Assertions.assertEquals(2, poolProxy.getTotalConnections());
Assertions.assertEquals(2, poolProxy.getActiveConnections());
Assertions.assertEquals(0, poolProxy.getIdleConnections());
Assertions.assertEquals(2, connectionPool.getTotalConnections());
Assertions.assertEquals(2, connectionPool.getActiveConnections());
Assertions.assertEquals(0, connectionPool.getIdleConnections());
}

private static HikariPoolMXBean getPoolProxy(Identity identity, ConnectionSpecification connectionSpecification, AuthenticationConfiguration authenticationConfiguration) throws MalformedObjectNameException
@Test
public void testConnectionPoolingForDifferentIdentities() throws Exception
{
String poolName = JDBCConnectionManager.getPoolName(identity, connectionSpecification, authenticationConfiguration);
MBeanServer mBeanServer = ManagementFactory.getPlatformMBeanServer();
return JMX.newMXBeanProxy(mBeanServer, new ObjectName("com.zaxxer.hikari:type=Pool (" + poolName + ")"), HikariPoolMXBean.class);
}
this.connectionFactory = new ConnectionFactory.Builder(this.environment, this.storeInstanceProvider)
.withCredentialBuilders(
new UserPasswordCredentialBuilder()
)
.withConnectionBuilders(
new StaticJDBCConnectionBuilder.WithPlaintextUsernamePassword()
)
.build();
this.storeInstanceProvider.injectStoreInstance(this.storeInstance);
Identity identity1 = identityFactory.createIdentity(
new IdentitySpecification.Builder()
.withName("testUser1")
.build()
);
Identity identity2 = identityFactory.createIdentity(
new IdentitySpecification.Builder()
.withName("testUser2")
.build()
);
ConnectionSpecification connectionSpecification = this.storeInstance.getConnectionSpecification();
AuthenticationConfiguration authenticationConfiguration = new UserPasswordAuthenticationConfiguration(
postgresContainer.getUser(),
new PropertiesFileSecret("passwordRef")
);

private static class InstrumentedStaticJDBCConnectionBuilder
{
static class WithPlaintextUsernamePassword extends StaticJDBCConnectionBuilder.WithPlaintextUsernamePassword
{
private final InstrumentedJDBCConnectionManager connectionManager;

WithPlaintextUsernamePassword(Function<HikariConfig, Void> hikariConfigHandler)
{
this.connectionManager = new InstrumentedJDBCConnectionManager(hikariConfigHandler);
}

@Override
public JDBCConnectionManager getConnectionManager()
{
return this.connectionManager;
}

@Override
protected Type[] actualTypeArguments()
{
Type genericSuperClass = this.getClass().getSuperclass().getGenericSuperclass();
ParameterizedType parameterizedType = (ParameterizedType) genericSuperClass;
return parameterizedType.getActualTypeArguments();
}
}
}
JDBCConnectionManager connectionManager = JDBCConnectionManager.getInstance();
Assertions.assertEquals(0, connectionManager.getPoolSize());

private static class InstrumentedJDBCConnectionManager extends JDBCConnectionManager
{
private final Function<HikariConfig, Void> hikariConfigHandler;
// 1. Get a new connection for identity1, which should initialize a pool
this.connectionFactory.getConnection(identity1, this.connectionFactory.getAuthenticator(identity1, TEST_STORE_INSTANCE_NAME, authenticationConfiguration));

InstrumentedJDBCConnectionManager(Function<HikariConfig, Void> hikariConfigHandler)
{
this.hikariConfigHandler = hikariConfigHandler;
}
String poolName1 = JDBCConnectionManager.getPoolName(identity1, connectionSpecification, authenticationConfiguration);
JDBCConnectionManager.ConnectionPool connectionPool1 = connectionManager.getPool(poolName1);

Assertions.assertEquals(1, connectionManager.getPoolSize());
Assertions.assertEquals(1, connectionPool1.getTotalConnections());
Assertions.assertEquals(1, connectionPool1.getActiveConnections());
Assertions.assertEquals(0, connectionPool1.getIdleConnections());

// 2. Get a new connection for identity2, which should initialize another pool
this.connectionFactory.getConnection(identity2, this.connectionFactory.getAuthenticator(identity2, TEST_STORE_INSTANCE_NAME, authenticationConfiguration));

@Override
protected void handleHikariConfig(HikariConfig config)
String poolName2 = JDBCConnectionManager.getPoolName(identity2, connectionSpecification, authenticationConfiguration);
JDBCConnectionManager.ConnectionPool connectionPool2 = connectionManager.getPool(poolName2);

Assertions.assertEquals(2, connectionManager.getPoolSize());
Assertions.assertEquals(1, connectionPool2.getTotalConnections());
Assertions.assertEquals(1, connectionPool2.getActiveConnections());
Assertions.assertEquals(0, connectionPool2.getIdleConnections());
}

public static class CustomAdvice
{
@Advice.OnMethodExit
public static void intercept(@Advice.Return(readOnly = false) String value)
{
config.setRegisterMbeans(true);
this.hikariConfigHandler.apply(config);
System.out.println("intercepted: " + value);
value = "hi: " + value;
}
}

// public static class MyWay
// {
// }
//
// private static class InstrumentedStaticJDBCConnectionBuilder
// {
// static class WithPlaintextUsernamePassword extends StaticJDBCConnectionBuilder.WithPlaintextUsernamePassword
// {
// private final InstrumentedJDBCConnectionManager connectionManager;
//
// WithPlaintextUsernamePassword(Function<HikariConfig, Void> hikariConfigHandler)
// {
// this.connectionManager = new InstrumentedJDBCConnectionManager(hikariConfigHandler);
// }
//
// @Override
// public JDBCConnectionManager getConnectionManager()
// {
// return this.connectionManager;
// }
//
// @Override
// protected Type[] actualTypeArguments()
// {
// Type genericSuperClass = this.getClass().getSuperclass().getGenericSuperclass();
// ParameterizedType parameterizedType = (ParameterizedType) genericSuperClass;
// return parameterizedType.getActualTypeArguments();
// }
// }
// }
//
// private static class InstrumentedJDBCConnectionManager extends JDBCConnectionManager
// {
// private final Function<HikariConfig, Void> hikariConfigHandler;
//
// InstrumentedJDBCConnectionManager(Function<HikariConfig, Void> hikariConfigHandler)
// {
// this.hikariConfigHandler = hikariConfigHandler;
// }
//
//// @Override
//// protected void handleHikariConfig(HikariConfig config)
//// {
//// config.setRegisterMbeans(true);
//// this.hikariConfigHandler.apply(config);
//// }
// }
}
4 changes: 2 additions & 2 deletions ...ication-connection-factory/src/main/java/org/finos/legend/connection/IdentityFactory.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,9 @@
import org.eclipse.collections.api.factory.Lists;
import org.finos.legend.engine.shared.core.identity.Credential;
import org.finos.legend.engine.shared.core.identity.Identity;
import org.finos.legend.engine.shared.core.identity.credential.AnonymousCredential;
import org.finos.legend.engine.shared.core.identity.factory.DefaultIdentityFactory;

import javax.security.auth.Subject;
import java.util.List;

public class IdentityFactory
Expand Down Expand Up @@ -49,7 +49,7 @@ public Identity createIdentity(IdentitySpecification identitySpecification)
}
if (credentials.isEmpty())
{
return DEFAULT.makeUnknownIdentity();
return identitySpecification.getName() != null ? new Identity(identitySpecification.getName(), new AnonymousCredential()) : DEFAULT.makeUnknownIdentity();
}
return new Identity(identitySpecification.getName(), credentials);
}
Expand Down
2 changes: 1 addition & 1 deletion ...n-connection-factory/src/main/java/org/finos/legend/connection/IdentitySpecification.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ public class IdentitySpecification

private IdentitySpecification(String name, List<CommonProfile> profiles, Subject subject, List<Credential> credentials)
{
this.name = Objects.requireNonNull(name, "Name is missing");
this.name = name;
this.profiles = profiles;
this.subject = subject;
this.credentials = credentials;
Expand Down
15 changes: 14 additions & 1 deletion ...onalStore-connection/src/main/java/org/finos/legend/connection/JDBCConnectionBuilder.java
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,19 @@
// Copyright 2023 Goldman Sachs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package org.finos.legend.connection;

import org.finos.legend.connection.ConnectionBuilder;
import org.finos.legend.connection.impl.JDBCConnectionManager;
import org.finos.legend.connection.protocol.ConnectionSpecification;
import org.finos.legend.engine.shared.core.identity.Credential;
Expand Down
Loading

0 comments on commit 7d2febd

Please sign in to comment.