Add affinity support

Author: rPraml

ebean-datasource-api/src/main/java/io/ebean/datasource/DataSourceBuilder.java

@@ -7,6 +7,7 @@
 import java.util.Properties;
 import java.util.function.BooleanSupplier;
 import java.util.function.Consumer;
+import java.util.function.Supplier;
 
 /**
  * Builder for DataSourcePool.
@@ -639,6 +640,26 @@ default DataSourceBuilder customProperties(Map<String, String> customProperties)
    */
   DataSourceBuilder addProperty(String key, int value);
 
+  /**
+   * sets the affinity-size (internal hashmap of distinct affinity keys). Should be a prime number. Default: 257
+   */
+  DataSourceBuilder affinitySize(int affinitySize);
+
+  /**
+   * Returns the affinity size.
+   */
+  int getAffinitySize();
+
+  /**
+   * Sets the affinity provider. e.g. Thread::currentThread.
+   */
+  DataSourceBuilder affinityProvider(Supplier<Object> affinityProvider);
+
+  /**
+   * Returns the affinity provider.
+   */
+  Supplier<Object> getAffinityProvider();
+
   /**
    * Set the database owner username (used to create connection for use with InitDatabase).
    */

ebean-datasource-api/src/main/java/io/ebean/datasource/DataSourceConfig.java

@@ -9,6 +9,7 @@
 import java.util.Map;
 import java.util.Properties;
 import java.util.function.Consumer;
+import java.util.function.Supplier;
 
 /**
  * Configuration information for a DataSource.
@@ -84,6 +85,8 @@ public class DataSourceConfig implements DataSourceBuilder.Settings {
   private boolean shutdownOnJvmExit;
   private boolean validateOnHeartbeat = !System.getenv().containsKey("LAMBDA_TASK_ROOT");
   private boolean enforceCleanClose;
+  private int affinitySize = 257;
+  private Supplier<Object> affinityProvider;
 
   @Override
   public Settings settings() {
@@ -146,6 +149,8 @@ public DataSourceConfig copy() {
     copy.alert = alert;
     copy.listener = listener;
     copy.enforceCleanClose = enforceCleanClose;
+    copy.affinitySize = affinitySize;
+    copy.affinityProvider = affinityProvider;
     return copy;
   }
 
@@ -654,6 +659,28 @@ public DataSourceConfig addProperty(String key, int value) {
     return addProperty(key, Integer.toString(value));
   }
 
+  @Override
+  public DataSourceBuilder affinitySize(int affinitySize) {
+    this.affinitySize = affinitySize;
+    return this;
+  }
+
+  @Override
+  public int getAffinitySize() {
+    return affinitySize;
+  }
+
+  @Override
+  public DataSourceBuilder affinityProvider(Supplier<Object> affinityProvider) {
+    this.affinityProvider = affinityProvider;
+    return this;
+  }
+
+  @Override
+  public Supplier<Object> getAffinityProvider() {
+    return affinityProvider;
+  }
+
   @Override
   public String getOwnerUsername() {
     return ownerUsername;
@@ -801,6 +828,7 @@ private void loadSettings(ConfigPropertiesHelper properties) {
     shutdownOnJvmExit = properties.getBoolean("shutdownOnJvmExit", shutdownOnJvmExit);
     validateOnHeartbeat = properties.getBoolean("validateOnHeartbeat", validateOnHeartbeat);
     enforceCleanClose = properties.getBoolean("enforceCleanClose", enforceCleanClose);
+    affinitySize = properties.getInt("affinityCacheSize", affinitySize);
 
 
     String isoLevel = properties.get("isolationLevel", _isolationLevel(isolationLevel));

ebean-datasource-api/src/main/java/io/ebean/datasource/DataSourceConnection.java

@@ -0,0 +1,17 @@
+package io.ebean.datasource;
+
+import java.sql.Connection;
+
+/**
+ * Interface for connection objects returned from the ebean-datasource connection pool
+ *
+ * @author Roland Praml, Foconis Analytics GmbH
+ */
+public interface DataSourceConnection extends Connection {
+
+  /**
+   * Returns the affinity-ID, this connection was assigned to.
+   */
+  Object affinityId();
+
+}

ebean-datasource-api/src/main/java/io/ebean/datasource/DataSourcePool.java

@@ -81,6 +81,12 @@ static DataSourceBuilder builder() {
    */
   void offline();
 
+  /**
+   * Returns a connection for given affinity ID. It is guaranteed, that connection.affinityId in listener etc.
+   * is the same object.
+   */
+  DataSourceConnection getConnection(Object affinityId) throws SQLException;
+
   /**
    * Shutdown the pool.
    * <p>