Add Additional dbcp2 configs

document-store/src/main/java/org/hypertrace/core/documentstore/model/config/ConnectionPoolConfig.java

@@ -1,14 +1,21 @@
 package org.hypertrace.core.documentstore.model.config;
 
 import java.time.Duration;
+import java.util.List;
 import javax.annotation.Nonnegative;
+import javax.annotation.Nullable;
 import lombok.AccessLevel;
 import lombok.AllArgsConstructor;
 import lombok.Builder;
 import lombok.NonNull;
+import lombok.Singular;
 import lombok.Value;
 import lombok.experimental.Accessors;
 
+/**
+ * See <a href="https://commons.apache.org/proper/commons-dbcp/configuration.html">DBCP2
+ * Configuration</a> for details
+ */
 @Value
 @Builder
 @Accessors(fluent = true)
@@ -29,6 +36,32 @@ public class ConnectionPoolConfig {
   // Time duration to wait for surrendering an idle connection back to the pool
   @NonNull @Builder.Default Duration connectionSurrenderTimeout = Duration.ofMinutes(5);
 
-  // Whether to validate connections when borrowing from the pool
-  @NonNull @Builder.Default Boolean testOnBorrow = true;
+  @Nullable String validationQuery;
+
+  @Nullable Duration validationQueryTimeout;
+
+  @Nullable Boolean testOnCreate;
+
+  // Setting this to false explictly since leaving it true by default has a non-trivial perf impact
+  @NonNull @Builder.Default Boolean testOnBorrow = false;
+
+  @Nullable Boolean testOnReturn;
+
+  @Nullable Boolean testWhileIdle;
+
+  @Nullable Duration timeBetweenEvictionRuns;
+
+  @Nullable Integer numTestsPerEvictionRun;
+
+  @Nullable Duration minEvictableIdleTime;
+
+  @Nullable Duration softMinEvictableIdleTime;
+
+  @Nullable Duration maxConnLifetime;
+
+  @NonNull
+  @Singular("connectionInitSql")
+  List<String> connectionInitSqls;
+
+  @Nullable Boolean lifo;
 }

document-store/src/main/java/org/hypertrace/core/documentstore/model/config/TypesafeConfigDatastoreConfigExtractor.java

@@ -35,6 +35,18 @@ public class TypesafeConfigDatastoreConfigExtractor {
   private static final String DEFAULT_MAX_IDLE_PERCENT_KEY = "maxIdlePercent";
   private static final String DEFAULT_MIN_IDLE_PERCENT_KEY = "minIdlePercent";
   private static final String DEFAULT_TEST_ON_BORROW_KEY = "testOnBorrow";
+  private static final String DEFAULT_VALIDATION_QUERY_KEY = "validationQuery";
+  private static final String DEFAULT_VALIDATION_QUERY_TIMEOUT_KEY = "validationQueryTimeout";
+  private static final String DEFAULT_TEST_ON_CREATE_KEY = "testOnCreate";
+  private static final String DEFAULT_TEST_ON_RETURN_KEY = "testOnReturn";
+  private static final String DEFAULT_TEST_WHILE_IDLE_KEY = "testWhileIdle";
+  private static final String DEFAULT_TIME_BETWEEN_EVICTION_RUNS_KEY = "timeBetweenEvictionRuns";
+  private static final String DEFAULT_NUM_TESTS_PER_EVICTION_RUN_KEY = "numTestsPerEvictionRun";
+  private static final String DEFAULT_MIN_EVICTABLE_IDLE_TIME_KEY = "minEvictableIdleTime";
+  private static final String DEFAULT_SOFT_MIN_EVICTABLE_IDLE_TIME_KEY = "softMinEvictableIdleTime";
+  private static final String DEFAULT_MAX_CONN_LIFETIME_KEY = "maxConnLifetime";
+  private static final String DEFAULT_CONNECTION_INIT_SQLS_KEY = "connectionInitSqls";
+  private static final String DEFAULT_LIFO_KEY = "lifo";
   private static final String DEFAULT_AGGREGATION_PIPELINE_MODE_KEY = "aggregationPipelineMode";
   private static final String DEFAULT_DATA_FRESHNESS_KEY = "dataFreshness";
   private static final String DEFAULT_QUERY_TIMEOUT_KEY = "queryTimeout";
@@ -90,6 +102,18 @@ private TypesafeConfigDatastoreConfigExtractor(
         .poolMaxIdlePercentKey(DEFAULT_MAX_IDLE_PERCENT_KEY)
         .poolMinIdlePercentKey(DEFAULT_MIN_IDLE_PERCENT_KEY)
         .poolTestOnBorrowKey(DEFAULT_TEST_ON_BORROW_KEY)
+        .poolValidationQueryKey(DEFAULT_VALIDATION_QUERY_KEY)
+        .poolValidationQueryTimeoutKey(DEFAULT_VALIDATION_QUERY_TIMEOUT_KEY)
+        .poolTestOnCreateKey(DEFAULT_TEST_ON_CREATE_KEY)
+        .poolTestOnReturnKey(DEFAULT_TEST_ON_RETURN_KEY)
+        .poolTestWhileIdleKey(DEFAULT_TEST_WHILE_IDLE_KEY)
+        .poolTimeBetweenEvictionRunsKey(DEFAULT_TIME_BETWEEN_EVICTION_RUNS_KEY)
+        .poolNumTestsPerEvictionRunKey(DEFAULT_NUM_TESTS_PER_EVICTION_RUN_KEY)
+        .poolMinEvictableIdleTimeKey(DEFAULT_MIN_EVICTABLE_IDLE_TIME_KEY)
+        .poolSoftMinEvictableIdleTimeKey(DEFAULT_SOFT_MIN_EVICTABLE_IDLE_TIME_KEY)
+        .poolMaxConnLifetimeKey(DEFAULT_MAX_CONN_LIFETIME_KEY)
+        .poolConnectionInitSqlsKey(DEFAULT_CONNECTION_INIT_SQLS_KEY)
+        .poolLifoKey(DEFAULT_LIFO_KEY)
         .aggregationPipelineMode(DEFAULT_AGGREGATION_PIPELINE_MODE_KEY)
         .dataFreshnessKey(DEFAULT_DATA_FRESHNESS_KEY)
         .queryTimeoutKey(DEFAULT_QUERY_TIMEOUT_KEY)
@@ -255,6 +279,96 @@ public TypesafeConfigDatastoreConfigExtractor poolTestOnBorrowKey(@NonNull final
     return this;
   }
 
+  public TypesafeConfigDatastoreConfigExtractor poolValidationQueryKey(@NonNull final String key) {
+    if (config.hasPath(key)) {
+      connectionPoolConfigBuilder.validationQuery(config.getString(key));
+    }
+    return this;
+  }
+
+  public TypesafeConfigDatastoreConfigExtractor poolValidationQueryTimeoutKey(
+      @NonNull final String key) {
+    if (config.hasPath(key)) {
+      connectionPoolConfigBuilder.validationQueryTimeout(config.getDuration(key));
+    }
+    return this;
+  }
+
+  public TypesafeConfigDatastoreConfigExtractor poolTestOnCreateKey(@NonNull final String key) {
+    if (config.hasPath(key)) {
+      connectionPoolConfigBuilder.testOnCreate(config.getBoolean(key));
+    }
+    return this;
+  }
+
+  public TypesafeConfigDatastoreConfigExtractor poolTestOnReturnKey(@NonNull final String key) {
+    if (config.hasPath(key)) {
+      connectionPoolConfigBuilder.testOnReturn(config.getBoolean(key));
+    }
+    return this;
+  }
+
+  public TypesafeConfigDatastoreConfigExtractor poolTestWhileIdleKey(@NonNull final String key) {
+    if (config.hasPath(key)) {
+      connectionPoolConfigBuilder.testWhileIdle(config.getBoolean(key));
+    }
+    return this;
+  }
+
+  public TypesafeConfigDatastoreConfigExtractor poolTimeBetweenEvictionRunsKey(
+      @NonNull final String key) {
+    if (config.hasPath(key)) {
+      connectionPoolConfigBuilder.timeBetweenEvictionRuns(config.getDuration(key));
+    }
+    return this;
+  }
+
+  public TypesafeConfigDatastoreConfigExtractor poolNumTestsPerEvictionRunKey(
+      @NonNull final String key) {
+    if (config.hasPath(key)) {
+      connectionPoolConfigBuilder.numTestsPerEvictionRun(config.getInt(key));
+    }
+    return this;
+  }
+
+  public TypesafeConfigDatastoreConfigExtractor poolMinEvictableIdleTimeKey(
+      @NonNull final String key) {
+    if (config.hasPath(key)) {
+      connectionPoolConfigBuilder.minEvictableIdleTime(config.getDuration(key));
+    }
+    return this;
+  }
+
+  public TypesafeConfigDatastoreConfigExtractor poolSoftMinEvictableIdleTimeKey(
+      @NonNull final String key) {
+    if (config.hasPath(key)) {
+      connectionPoolConfigBuilder.softMinEvictableIdleTime(config.getDuration(key));
+    }
+    return this;
+  }
+
+  public TypesafeConfigDatastoreConfigExtractor poolMaxConnLifetimeKey(@NonNull final String key) {
+    if (config.hasPath(key)) {
+      connectionPoolConfigBuilder.maxConnLifetime(config.getDuration(key));
+    }
+    return this;
+  }
+
+  public TypesafeConfigDatastoreConfigExtractor poolConnectionInitSqlsKey(
+      @NonNull final String key) {
+    if (config.hasPath(key)) {
+      config.getStringList(key).forEach(connectionPoolConfigBuilder::connectionInitSql);
+    }
+    return this;
+  }
+
+  public TypesafeConfigDatastoreConfigExtractor poolLifoKey(@NonNull final String key) {
+    if (config.hasPath(key)) {
+      connectionPoolConfigBuilder.lifo(config.getBoolean(key));
+    }
+    return this;
+  }
+
   public TypesafeConfigDatastoreConfigExtractor aggregationPipelineMode(@NonNull final String key) {
     if (config.hasPath(key)) {
       connectionConfigBuilder.aggregationPipelineMode(

document-store/src/main/java/org/hypertrace/core/documentstore/model/config/postgres/PostgresConnectionConfig.java

@@ -17,19 +17,22 @@
 import lombok.Value;
 import lombok.experimental.Accessors;
 import lombok.experimental.NonFinal;
+import lombok.extern.slf4j.Slf4j;
 import org.hypertrace.core.documentstore.model.config.ConnectionConfig;
 import org.hypertrace.core.documentstore.model.config.ConnectionCredentials;
 import org.hypertrace.core.documentstore.model.config.ConnectionPoolConfig;
 import org.hypertrace.core.documentstore.model.config.DatabaseType;
 import org.hypertrace.core.documentstore.model.config.Endpoint;
 import org.postgresql.PGProperty;
 
+@Slf4j
 @Value
 @NonFinal
 @Accessors(fluent = true)
 @ToString(callSuper = true)
 @EqualsAndHashCode(callSuper = true)
 public class PostgresConnectionConfig extends ConnectionConfig {
+
   private static final ConnectionCredentials DEFAULT_CREDENTIALS =
       ConnectionCredentials.builder()
           .username(PostgresDefaults.DEFAULT_USER)
@@ -41,12 +44,14 @@ public class PostgresConnectionConfig extends ConnectionConfig {
 
   private static final String SCHEMA_CACHE_EXPIRY_MS_KEY = "schemaCacheExpiryMs";
   private static final String SCHEMA_REFRESH_COOLDOWN_MS_KEY = "schemaRefreshCooldownMs";
+  private static final String MIN_SERVER_VERSION_KEY = "minServerVersion";
 
   @NonNull String applicationName;
   @NonNull ConnectionPoolConfig connectionPoolConfig;
   @NonNull Duration queryTimeout;
   @NonNull Duration schemaCacheExpiry;
   @NonNull Duration schemaRefreshCooldown;
+  @Nullable String minServerVersion;
   @NonNull Map<String, CollectionConfig> collectionConfigs;
 
   public static ConnectionConfigBuilder builder() {
@@ -72,6 +77,7 @@ public PostgresConnectionConfig(
     this.queryTimeout = queryTimeout;
     this.schemaCacheExpiry = extractSchemaCacheExpiry(customParameters);
     this.schemaRefreshCooldown = extractSchemaRefreshCooldown(customParameters);
+    this.minServerVersion = extractMinServerVersion(customParameters);
     this.collectionConfigs = collectionConfigs != null ? collectionConfigs : Collections.emptyMap();
   }
 
@@ -91,6 +97,17 @@ public Properties buildProperties() {
     }
 
     properties.setProperty(PGProperty.APPLICATION_NAME.getName(), applicationName());
+    Optional.ofNullable(minServerVersion)
+        .ifPresent(
+            version ->
+                properties.setProperty(PGProperty.ASSUME_MIN_SERVER_VERSION.getName(), version));
+
+    log.debug(
+        "Postgres JDBC properties - {}={}, {}={}",
+        PGProperty.APPLICATION_NAME,
+        properties.getProperty(PGProperty.APPLICATION_NAME.getName()),
+        PGProperty.ASSUME_MIN_SERVER_VERSION,
+        properties.getProperty(PGProperty.ASSUME_MIN_SERVER_VERSION.getName()));
 
     return properties;
   }
@@ -162,6 +179,13 @@ private static Duration extractSchemaRefreshCooldown(final Map<String, String> c
         .orElse(DEFAULT_SCHEMA_REFRESH_COOLDOWN);
   }
 
+  @Nullable
+  private static String extractMinServerVersion(final Map<String, String> customParameters) {
+    return Optional.ofNullable(customParameters.get(MIN_SERVER_VERSION_KEY))
+        .filter(not(String::isBlank))
+        .orElse(null);
+  }
+
   public Optional<CollectionConfig> getCollectionConfig(String collectionName) {
     return Optional.ofNullable(collectionConfigs.get(collectionName));
   }

document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresConnectionPool.java

@@ -9,17 +9,13 @@
 import org.apache.commons.dbcp2.PoolableConnection;
 import org.apache.commons.dbcp2.PoolableConnectionFactory;
 import org.apache.commons.dbcp2.PoolingDataSource;
-import org.apache.commons.pool2.impl.AbandonedConfig;
 import org.apache.commons.pool2.impl.GenericObjectPool;
 import org.hypertrace.core.documentstore.model.config.ConnectionPoolConfig;
 import org.hypertrace.core.documentstore.model.config.postgres.PostgresConnectionConfig;
 
 @Slf4j
 class PostgresConnectionPool {
 
-  private static final String VALIDATION_QUERY = "SELECT 1";
-  private static final Duration VALIDATION_QUERY_TIMEOUT = Duration.ofSeconds(5);
-
   // Single pool with autoCommit=true by default. Transactional connections flip autoCommit
   // to false on borrow; DBCP2 resets it back when the connection is returned to the pool.
   private final PoolingDataSource<PoolableConnection> dataSource;
@@ -66,57 +62,88 @@ private PoolingDataSource<PoolableConnection> createPooledDataSource(
 
     final ConnectionPoolConfig poolConfig = config.connectionPoolConfig();
     setPoolProperties(connectionPool, poolConfig);
-    setFactoryProperties(poolableConnectionFactory, connectionPool);
+    setFactoryProperties(poolableConnectionFactory, connectionPool, poolConfig);
 
     return new PoolingDataSource<>(connectionPool);
   }
 
   private void setPoolProperties(
       final GenericObjectPool<PoolableConnection> connectionPool,
       final ConnectionPoolConfig poolConfig) {
-    final AbandonedConfig abandonedConfig = getAbandonedConfig(poolConfig);
     final int maxConnections = poolConfig.maxConnections();
     connectionPool.setMaxTotal(maxConnections);
     connectionPool.setMaxIdle(getIdleCount(poolConfig.maxIdlePercent(), maxConnections));
     connectionPool.setMinIdle(getIdleCount(poolConfig.minIdlePercent(), maxConnections));
     connectionPool.setBlockWhenExhausted(true);
     connectionPool.setMaxWaitMillis(poolConfig.connectionAccessTimeout().toMillis());
+    if (poolConfig.testOnCreate() != null) {
+      connectionPool.setTestOnCreate(poolConfig.testOnCreate());
+    }
     connectionPool.setTestOnBorrow(poolConfig.testOnBorrow());
-    connectionPool.setAbandonedConfig(abandonedConfig);
+    if (poolConfig.testOnReturn() != null) {
+      connectionPool.setTestOnReturn(poolConfig.testOnReturn());
+    }
+    if (poolConfig.testWhileIdle() != null) {
+      connectionPool.setTestWhileIdle(poolConfig.testWhileIdle());
+    }
+    if (poolConfig.timeBetweenEvictionRuns() != null) {
+      connectionPool.setTimeBetweenEvictionRuns(poolConfig.timeBetweenEvictionRuns());
+    }
+    if (poolConfig.numTestsPerEvictionRun() != null) {
+      connectionPool.setNumTestsPerEvictionRun(poolConfig.numTestsPerEvictionRun());
+    }
+    if (poolConfig.minEvictableIdleTime() != null) {
+      connectionPool.setMinEvictableIdleTime(poolConfig.minEvictableIdleTime());
+    }
+    if (poolConfig.softMinEvictableIdleTime() != null) {
+      connectionPool.setSoftMinEvictableIdleTime(poolConfig.softMinEvictableIdleTime());
+    }
+    if (poolConfig.lifo() != null) {
+      connectionPool.setLifo(poolConfig.lifo());
+    }
     log.debug(
-        "Postgres connection pool properties - maxTotal: {}, maxIdle: {}, minIdle: {}, maxWaitMillis: {}, connectionSurrenderTimeout: {}, testOnBorrow: {}",
+        "Postgres connection pool properties - maxTotal: {}, maxIdle: {}, minIdle: {}, maxWaitMillis: {}, testOnBorrow: {}, testWhileIdle: {}, timeBetweenEvictionRunsMillis: {}, minEvictableIdleTimeMillis: {}, maxConnLifetimeMillis: {}, lifo: {}",
         connectionPool.getMaxTotal(),
         connectionPool.getMaxIdle(),
         connectionPool.getMinIdle(),
         connectionPool.getMaxWaitMillis(),
-        poolConfig.connectionSurrenderTimeout(),
-        poolConfig.testOnBorrow());
+        poolConfig.testOnBorrow(),
+        poolConfig.testWhileIdle(),
+        toMillisOrNull(poolConfig.timeBetweenEvictionRuns()),
+        toMillisOrNull(poolConfig.minEvictableIdleTime()),
+        toMillisOrNull(poolConfig.maxConnLifetime()),
+        poolConfig.lifo());
+  }
+
+  private static Long toMillisOrNull(final Duration duration) {
+    return duration == null ? null : duration.toMillis();
   }
 
   private void setFactoryProperties(
       PoolableConnectionFactory poolableConnectionFactory,
-      GenericObjectPool<PoolableConnection> connectionPool) {
+      GenericObjectPool<PoolableConnection> connectionPool,
+      ConnectionPoolConfig poolConfig) {
     poolableConnectionFactory.setPool(connectionPool);
-    poolableConnectionFactory.setValidationQuery(VALIDATION_QUERY);
-    poolableConnectionFactory.setValidationQueryTimeout((int) VALIDATION_QUERY_TIMEOUT.toSeconds());
+    if (poolConfig.validationQuery() != null) {
+      poolableConnectionFactory.setValidationQuery(poolConfig.validationQuery());
+    }
+    if (poolConfig.validationQueryTimeout() != null) {
+      poolableConnectionFactory.setValidationQueryTimeout(
+          (int) poolConfig.validationQueryTimeout().toSeconds());
+    }
     poolableConnectionFactory.setDefaultReadOnly(false);
     poolableConnectionFactory.setDefaultAutoCommit(true);
+    if (poolConfig.maxConnLifetime() != null) {
+      poolableConnectionFactory.setMaxConnLifetimeMillis(poolConfig.maxConnLifetime().toMillis());
+    }
+    poolableConnectionFactory.setConnectionInitSql(poolConfig.connectionInitSqls());
     // Note: We intentionally do NOT call setDefaultTransactionIsolation() here.
     // PostgreSQL defaults to READ_COMMITTED, which is what we want. Setting it explicitly
     // causes DBCP2 to execute "SHOW TRANSACTION ISOLATION LEVEL" on every connection borrow
     // (via PgConnection.getTransactionIsolation()), adding unnecessary overhead.
     poolableConnectionFactory.setPoolStatements(false);
   }
 
-  private AbandonedConfig getAbandonedConfig(final ConnectionPoolConfig poolConfig) {
-    final AbandonedConfig abandonedConfig = new AbandonedConfig();
-    abandonedConfig.setLogAbandoned(true);
-    abandonedConfig.setRemoveAbandonedOnBorrow(true);
-    abandonedConfig.setRequireFullStackTrace(true);
-    abandonedConfig.setRemoveAbandonedTimeout(poolConfig.connectionSurrenderTimeout());
-    return abandonedConfig;
-  }
-
   private int getIdleCount(final int percent, final int maxConnections) {
     if (percent < 0) {
       return maxConnections;