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

@@ -36,7 +36,6 @@ public interface DataSourcePool extends DataSource {
    *   Connection connection = pool.getConnection();
    *
    * }</pre>
-   *
    */
   static DataSourceBuilder builder() {
     return new DataSourceConfig();
@@ -81,6 +80,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>

ebean-datasource/src/main/java/io/ebean/datasource/pool/ConnectionBuffer.java

@@ -11,14 +11,107 @@
  * PooledConnection, so that the node object can be reused. This avoids object
  * creation/gc during remove operations.
  * <p>
+ * The connectionbuffer iself has one linkedList from <code>free</code> to
+ * <code>freeEnd</code>. In parallel, the elements in this list can also be part
+ * the affinityNodes list, which implement a kind of hashmap.
+ * <p>
+ * So you can prefer which connection should be taken. You can use CurrentThread or
+ * currentTenant as affinity ID. So you likely get a connection that has the right
+ * pstatement caches or is already in the CPU cache.
+ * <p>
+ * Without affinityId, the first free-connection is taken.
+ * <p>
+ * With affinityId, the affinityNodes-list is determined by the hashCode, then the
+ * list is searched, if there is a connection with the same affinity object.
+ * <p>
+ * If there is no one found, we take the LAST connection in freeList, as this is
+ * the best candidate not to steal the affinity of a connection, that was currently
+ * used. This ensures (or also causes) that the pool has at least that size of the
+ * frequent used affinityIds. E.g. if the affinity id represents tenant id, and
+ * 15 tenants are active, the pool will not shrink below 15 - on the other hand,
+ * there is always one connection ready for each active tenant.
+ * <p>
+ * A free node can be member in two lists:
+ * <ol>
+ *     <li>it is definitively member in the freeList</li>
+ *     <li>it may be member in one of the affinity-lists (mod hash)</li>
+ * </ol>
+ * The remove / transition from free to busy will remove the node from both lists.
+ * <p>
+ * Graphical exammple
+ * <pre>
+ *     By default, the busy list is empty
+ *     busy ---------------------------------------------------> busyEnd
+ *     free --> c1 --> c2 --> c3 --> c4 --> c5 --> c6 --> c7 --> freeEnd
+ *     al1  ---------------------------------------------------> end
+ *     al2  ---------------------------------------------------> end
+ *     ...
+ *     al257---------------------------------------------------> end
+ *
+ *     if a popFree(1) is called, we lookup in al1 and found no usable node.
+ *     in this case, we take the last node, c7 and move it to the busy list
+ *
+ *     busy --> c7 --------------------------------------------> busyEnd
+ *     free --> c1 --> c2 --> c3 --> c4 --> c5 --> c6 ---------> freeEnd
+ *
+ *     When we put that node back in the freelist, it becomes the first node
+ *     and it will be also linked in affinity-list1
+ *
+ *     busy ---------------------------------------------------> busyEnd
+ *     free ,       ,> c1 --> c2 --> c3 --> c4 --> c5 --> c6 --> freeEnd
+ *     al1-> '> c7 '
+ *     al2-> (empty)
+ *
+ *     subsequent popFree(1) will always return c7 as long as it is not busy.
+ *     now we call popFree(1) twice, we will get this picture
+ *
+ *     busy --> c6 --> c7 ----------------------------------------> busyEnd
+ *     free --> c1 --> c2 --> c3 --> c4 --> c5 -------------------> freeEnd
+ *     al1-> (empty)
+ *     al2-> (empty)
+ *
+ *     putting them back
+ *
+ *     busy ------------------------------------------------------> busyEnd
+ *     free ,             ,> c1 --> c2 --> c3 --> c4 --> c5 -----> freeEnd
+ *     al1-> '> c7 -- c6 '
+ *     al2-> (empty)
+ *
+ *     fetching and return a connection with affinity = 2:
+ *
+ *     busy ------------------------------------------------------> busyEnd
+ *     free ,                   ,> c1 --> c2 --> c3 --> c4 -------> freeEnd
+ *     al1-> |     '> c7 -- c6 '
+ *     al2-> '> c2 '
+ *
+ *     so we have 2 connections for affinity 1 and one connection for affinity 2
+ *     (and the rest is ordered itself in the freeList)
+ * </pre>
+ * <p>
  * All thread safety controlled externally (by PooledConnectionQueue).
  * </p>
  */
 final class ConnectionBuffer {
 
     private final Node free = Node.init();
+    private final Node freeEnd = free.next;
     private final Node busy = Node.init();
 
+    private final Node[] affinityNodes;
+    private final int hashSize;
+
+    ConnectionBuffer(int hashSize) {
+        this.hashSize = hashSize;
+        if (hashSize > 0) {
+            affinityNodes = new Node[hashSize];
+            for (int i = 0; i < affinityNodes.length; i++) {
+                affinityNodes[i] = Node.init();
+            }
+        } else {
+            affinityNodes = null;
+        }
+    }
+
     int freeSize = 0;
     int busySize = 0;
 
@@ -76,17 +169,36 @@ boolean moveToFreeList(PooledConnection c) {
         }
         node.remove();
         busySize--;
-        node.addAfter(free);
+        Object affinityId = c.affinityId();
+        if (affinityId != null) {
+            node.addAfter(free, affinityNodes[affinityId.hashCode() % hashSize]);
+        } else {
+            node.addAfter(free);
+        }
         freeSize++;
         c.setBusyNode(null);
         return true;
     }
 
     /**
      * Remove a connection from the free list. Returns <code>null</code> if there is not any.
+     * <p>
+     * Connections that are returend from this method must be either added to busyList with
+     * addBusy or closed fully.
      */
-    PooledConnection popFree() {
-        Node node = free.next;
+    PooledConnection popFree(Object affinityId) {
+        Node node;
+        if (affinityId == null || affinityNodes == null) {
+            node = free.next;
+        } else {
+            node = affinityNodes[affinityId.hashCode() % hashSize].find(affinityId);
+            if (node == null) {
+                // when we did not find a node with that affinity, we take the last (oldest one)
+                // and reuse this with the new affinity. This avoids to "steal" the affinity
+                // from the newest one.
+                node = freeEnd.prev;
+            }
+        }
         if (node.isBoundaryNode()) {
             return null;
         }
@@ -115,10 +227,10 @@ int addBusy(PooledConnection c) {
      */
     void closeAllFree(boolean logErrors) {
         List<PooledConnection> tempList = new ArrayList<>();
-        PooledConnection c = popFree();
+        PooledConnection c = popFree(null);
         while (c != null) {
             tempList.add(c);
-            c = popFree();
+            c = popFree(null);
         }
 
         if (Log.isLoggable(System.Logger.Level.TRACE)) {
@@ -212,20 +324,25 @@ static final class Node {
 
         private Node next;
         private Node prev;
+        // Double-LL nodes for affinity management
+        private Node afNext;
+        private Node afPrev;
         final PooledConnection pc;
 
         private Node(PooledConnection pc) {
             this.pc = pc;
         }
 
         /**
-         * Creates new "list" with two empty boundary nodes
+         * Creates new "list" with two empty edge nodes
          */
         public static Node init() {
             Node node1 = new Node(null);
             Node node2 = new Node(null);
             node1.next = node2;
             node2.prev = node1;
+            node1.afNext = node2;
+            node2.afPrev = node1;
             return node1;
         }
 
@@ -240,12 +357,18 @@ private boolean isBoundaryNode() {
          * Removes the node from the list. The node can be re-added to an other list
          */
         private void remove() {
-            assert pc != null : "called remove a boundary node";
+            assert pc != null : "called remove on an edge node";
             assert prev != null && next != null : "not part of a list";
             next.prev = prev;
             prev.next = next;
             prev = null;
             next = null;
+            if (afNext != null) {
+                afNext.afPrev = afPrev;
+                afPrev.afNext = afNext;
+                afPrev = null;
+                afNext = null;
+            }
         }
 
         /**
@@ -255,11 +378,37 @@ private void remove() {
          */
         public void addAfter(Node node) {
             assert !this.isBoundaryNode() : "this is a boundary node";
-            assert next == null & prev == null : "Node already member of a list";
+            assert next == null && prev == null : "Node already member of a list";
             next = node.next;
             prev = node;
             node.next.prev = this;
             node.next = this;
         }
+
+        /**
+         * Adds <code>this</code> after <code>node</code> AND as affinity-node after <code>afNode</code>.
+         */
+        public void addAfter(Node node, Node afNode) {
+            addAfter(node);
+            assert afNext == null && afPrev == null : "Node already member of affinity-list";
+            afNext = afNode.afNext;
+            afPrev = afNode;
+            afNode.afNext.afPrev = this;
+            afNode.afNext = this;
+        }
+
+        /**
+         * Find the connection with given affinity id in this affinity-list.
+         */
+        public Node find(Object affinityId) {
+            Node n = this.afNext;
+            while (!n.isBoundaryNode()) {
+                if (affinityId.equals(n.pc.affinityId())) {
+                    return n;
+                }
+                n = n.afNext;
+            }
+            return null;
+        }
     }
 }