IMPROVED: Tenant support for query plans

ebean-api/src/main/java/io/ebean/meta/MetaQueryPlan.java

@@ -44,6 +44,11 @@ public interface MetaQueryPlan {
    */
   String plan();
 
+  /**
+   * The tenant ID of the plan.
+   */
+  Object tenantId();
+
   /**
    * Return the query execution time associated with the bind values capture.
    */

ebean-core/src/main/java/io/ebeaninternal/api/SpiDbQueryPlan.java

@@ -12,6 +12,6 @@ public interface SpiDbQueryPlan extends MetaQueryPlan {
   /**
    * Extend with queryTimeMicros, captureCount, captureMicros and when the bind values were captured.
    */
-  SpiDbQueryPlan with(long queryTimeMicros, long captureCount, long captureMicros, Instant whenCaptured);
+  SpiDbQueryPlan with(long queryTimeMicros, long captureCount, long captureMicros, Instant whenCaptured, Object tenantId);
 
 }

ebean-core/src/main/java/io/ebeaninternal/api/SpiTransactionManager.java

@@ -60,6 +60,6 @@ public interface SpiTransactionManager {
   /**
    * Return a connection used for query plan collection.
    */
-  Connection queryPlanConnection() throws SQLException;
+  Connection queryPlanConnection(Object tenantId) throws SQLException;
 
 }

ebean-core/src/main/java/io/ebeaninternal/server/core/InternalConfiguration.java

@@ -583,7 +583,8 @@ public QueryPlanManager initQueryPlanManager(TransactionManager transactionManag
       return QueryPlanManager.NOOP;
     }
     long threshold = config.getQueryPlanThresholdMicros();
-    return new CQueryPlanManager(transactionManager, threshold, queryPlanLogger(databasePlatform.platform()), extraMetrics);
+    return new CQueryPlanManager(transactionManager, config.getCurrentTenantProvider(),
+      threshold, queryPlanLogger(databasePlatform.platform()), extraMetrics);
   }
 
   /**

ebean-core/src/main/java/io/ebeaninternal/server/query/CQueryBindCapture.java

@@ -1,33 +1,39 @@
 package io.ebeaninternal.server.query;
 
-import io.ebeaninternal.api.SpiDbQueryPlan;
-import io.ebeaninternal.api.SpiQueryBindCapture;
-import io.ebeaninternal.api.SpiQueryPlan;
+import io.ebean.config.CurrentTenantProvider;
+import io.ebeaninternal.api.*;
 import io.ebeaninternal.server.bind.capture.BindCapture;
 
+import java.sql.Connection;
+import java.sql.SQLException;
 import java.time.Instant;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.ReentrantLock;
 
+import static java.lang.System.Logger.Level.ERROR;
+
 final class CQueryBindCapture implements SpiQueryBindCapture {
 
   private static final double multiplier = 1.5d;
 
   private final ReentrantLock lock = new ReentrantLock();
   private final CQueryPlanManager manager;
   private final SpiQueryPlan queryPlan;
+  private final CurrentTenantProvider tenantProvider;
 
   private BindCapture bindCapture;
   private long queryTimeMicros;
   private long thresholdMicros;
   private long captureCount;
+  private Object tenantId;
 
   private long lastBindCapture;
 
-  CQueryBindCapture(CQueryPlanManager manager, SpiQueryPlan queryPlan, long thresholdMicros) {
+  CQueryBindCapture(CQueryPlanManager manager, SpiQueryPlan queryPlan, long thresholdMicros, CurrentTenantProvider tenantProvider) {
     this.manager = manager;
     this.queryPlan = queryPlan;
     this.thresholdMicros = thresholdMicros;
+    this.tenantProvider = tenantProvider;
   }
 
   /**
@@ -45,6 +51,7 @@ public void setBind(BindCapture bindCapture, long queryTimeMicros, long startNan
       this.thresholdMicros = Math.round(queryTimeMicros * multiplier);
       this.captureCount++;
       this.bindCapture = bindCapture;
+      this.tenantId = tenantProvider == null ? null : tenantProvider.currentId();
       this.queryTimeMicros = queryTimeMicros;
       lastBindCapture = System.currentTimeMillis();
       manager.notifyBindCapture(this, startNanos);
@@ -68,19 +75,27 @@ public void queryPlanInit(long thresholdMicros) {
   /**
    * Collect the query plan using already captured bind values.
    */
-  public boolean collectQueryPlan(CQueryPlanRequest request) {
+  public boolean collectQueryPlan(CQueryPlanRequest request, SpiTransactionManager transactionManager) {
     if (bindCapture == null || request.since() < lastBindCapture) {
       // no bind capture since the last capture
       return false;
     }
 
+
     final Instant whenCaptured = Instant.ofEpochMilli(this.lastBindCapture);
     final BindCapture last = this.bindCapture;
+    final Object tenantId = this.tenantId;
     final long startNanos = System.nanoTime();
-    SpiDbQueryPlan queryPlan = manager.collectPlan(request.connection(), this.queryPlan, last);
+    SpiDbQueryPlan queryPlan;
+    try (Connection connection = transactionManager.queryPlanConnection(tenantId)) {
+      queryPlan = manager.collectPlan(connection, this.queryPlan, last);
+    } catch (SQLException e) {
+      CoreLog.log.log(ERROR, "Error during query plan collection", e);
+      return false;
+    }
     if (queryPlan != null) {
       final long captureMicros = TimeUnit.MICROSECONDS.convert(System.nanoTime() - startNanos, TimeUnit.NANOSECONDS);
-      request.add(queryPlan.with(queryTimeMicros, captureCount, captureMicros, whenCaptured));
+      request.add(queryPlan.with(queryTimeMicros, captureCount, captureMicros, whenCaptured, tenantId));
       // effectively turn off bind capture for this plan
       thresholdMicros = Long.MAX_VALUE;
       return true;

ebean-core/src/main/java/io/ebeaninternal/server/query/CQueryPlanManager.java

@@ -1,5 +1,6 @@
 package io.ebeaninternal.server.query;
 
+import io.ebean.config.CurrentTenantProvider;
 import io.ebean.meta.MetaQueryPlan;
 import io.ebean.meta.QueryPlanRequest;
 import io.ebean.metric.TimedMetric;
@@ -8,11 +9,9 @@
 import io.ebeaninternal.server.bind.capture.BindCapture;
 
 import java.sql.Connection;
-import java.sql.SQLException;
 import java.util.List;
 import java.util.concurrent.ConcurrentHashMap;
 
-import static java.lang.System.Logger.Level.ERROR;
 import static java.util.Collections.emptyList;
 
 public final class CQueryPlanManager implements QueryPlanManager {
@@ -21,13 +20,17 @@ public final class CQueryPlanManager implements QueryPlanManager {
 
   private final ConcurrentHashMap<CQueryBindCapture, Object> plans = new ConcurrentHashMap<>();
   private final TransactionManager transactionManager;
+  private final CurrentTenantProvider tenantProvider;
   private final QueryPlanLogger planLogger;
   private final TimedMetric timeCollection;
   private final TimedMetric timeBindCapture;
   private long defaultThreshold;
 
-  public CQueryPlanManager(TransactionManager transactionManager, long defaultThreshold, QueryPlanLogger planLogger, ExtraMetrics extraMetrics) {
+  public CQueryPlanManager(TransactionManager transactionManager,
+                           CurrentTenantProvider tenantProvider,
+                           long defaultThreshold, QueryPlanLogger planLogger, ExtraMetrics extraMetrics) {
     this.transactionManager = transactionManager;
+    this.tenantProvider = tenantProvider;
     this.defaultThreshold = defaultThreshold;
     this.planLogger = planLogger;
     this.timeCollection = extraMetrics.planCollect();
@@ -41,7 +44,7 @@ public void setDefaultThreshold(long thresholdMicros) {
 
   @Override
   public SpiQueryBindCapture createBindCapture(SpiQueryPlan queryPlan) {
-    return new CQueryBindCapture(this, queryPlan, defaultThreshold);
+    return new CQueryBindCapture(this, queryPlan, defaultThreshold, tenantProvider);
   }
 
   public void notifyBindCapture(CQueryBindCapture planBind, long startNanos) {
@@ -58,16 +61,12 @@ public List<MetaQueryPlan> collect(QueryPlanRequest request) {
   }
 
   private List<MetaQueryPlan> collectPlans(QueryPlanRequest request) {
-    try (Connection connection = transactionManager.queryPlanConnection()) {
-      CQueryPlanRequest req = new CQueryPlanRequest(connection, request, plans.keySet().iterator());
+
+    CQueryPlanRequest req = new CQueryPlanRequest(transactionManager, request, plans.keySet().iterator());
       while (req.hasNext()) {
         req.nextCapture();
       }
       return req.plans();
-    } catch (SQLException e) {
-      CoreLog.log.log(ERROR, "Error during query plan collection", e);
-      return emptyList();
-    }
   }
 
   public SpiDbQueryPlan collectPlan(Connection connection, SpiQueryPlan queryPlan, BindCapture last) {

ebean-core/src/main/java/io/ebeaninternal/server/query/CQueryPlanRequest.java

@@ -2,8 +2,8 @@
 
 import io.ebean.meta.MetaQueryPlan;
 import io.ebean.meta.QueryPlanRequest;
+import io.ebeaninternal.api.SpiTransactionManager;
 
-import java.sql.Connection;
 import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.List;
@@ -15,14 +15,15 @@ final class CQueryPlanRequest {
 
   private final List<MetaQueryPlan> plans = new ArrayList<>();
 
-  private final Connection connection;
+  private final SpiTransactionManager transactionManager;
   private final long since;
   private final int maxCount;
   private final long maxTime;
   private final Iterator<CQueryBindCapture> iterator;
 
-  CQueryPlanRequest(Connection connection, QueryPlanRequest req, Iterator<CQueryBindCapture> iterator) {
-    this.connection = connection;
+
+  CQueryPlanRequest(SpiTransactionManager transactionManager, QueryPlanRequest req, Iterator<CQueryBindCapture> iterator) {
+    this.transactionManager = transactionManager;
     this.iterator = iterator;
     this.maxCount = req.maxCount();
     long reqSince = req.since();
@@ -31,13 +32,6 @@ final class CQueryPlanRequest {
     this.maxTime = maxTimeMillis > 0 ? System.currentTimeMillis() + maxTimeMillis : 0;
   }
 
-  /**
-   * Return the connection used to collect the db query plan.
-   */
-  Connection connection() {
-    return connection;
-  }
-
   /**
    * Add the collected query plan.
    */
@@ -71,7 +65,7 @@ boolean hasNext() {
    */
   void nextCapture() {
     final CQueryBindCapture next = iterator.next();
-    if (next.collectQueryPlan(this)) {
+    if (next.collectQueryPlan(this, transactionManager)) {
       iterator.remove();
     }
   }

ebean-core/src/main/java/io/ebeaninternal/server/query/DQueryPlanOutput.java

@@ -24,6 +24,8 @@ final class DQueryPlanOutput implements MetaQueryPlan, SpiDbQueryPlan {
   private long captureMicros;
   private Instant whenCaptured;
 
+  private Object tenantId;
+
   DQueryPlanOutput(Class<?> beanType, String label, String hash, String sql, ProfileLocation profileLocation, String bind, String plan) {
     this.beanType = beanType;
     this.label = label;
@@ -84,6 +86,14 @@ public String plan() {
     return plan;
   }
 
+  /**
+   * Returns the tenant id of this plan.
+   */
+  @Override
+  public Object tenantId() {
+    return tenantId;
+  }
+
   /**
    * Return the query execution time associated with the capture of bind values used
    * to build the query plan.
@@ -113,18 +123,27 @@ public Instant whenCaptured() {
 
   @Override
   public String toString() {
-    return " BeanType:" + ((beanType == null) ? "" : beanType.getSimpleName()) + " planHash:" + hash + " label:" + label + " queryTimeMicros:" + queryTimeMicros + " captureCount:" + captureCount + "\n SQL:" + sql + "\nBIND:" + bind + "\nPLAN:" + plan;
+    return " BeanType:" + ((beanType == null) ? "" : beanType.getSimpleName())
+      + " planHash:" + hash
+      + " label:" + label
+      + " queryTimeMicros:" + queryTimeMicros
+      + " captureCount:" + captureCount
+      + (tenantId == null ? "" : (" tenant:" + tenantId))
+      + "\n SQL:" + sql
+      + "\nBIND:" + bind
+      + "\nPLAN:" + plan;
   }
 
   /**
    * Additionally set the query execution time and the number of bind captures.
    */
   @Override
-  public DQueryPlanOutput with(long queryTimeMicros, long captureCount, long captureMicros, Instant whenCaptured) {
+  public DQueryPlanOutput with(long queryTimeMicros, long captureCount, long captureMicros, Instant whenCaptured, Object tenantId) {
     this.queryTimeMicros = queryTimeMicros;
     this.captureCount = captureCount;
     this.captureMicros = captureMicros;
     this.whenCaptured = whenCaptured;
+    this.tenantId = tenantId;
     return this;
   }
 }

ebean-core/src/main/java/io/ebeaninternal/server/transaction/TransactionManager.java

@@ -247,8 +247,8 @@ public final String name() {
   }
 
   @Override
-  public final Connection queryPlanConnection() throws SQLException {
-    return dataSourceSupplier.connection(null);
+  public final Connection queryPlanConnection(Object tenantId) throws SQLException {
+    return dataSourceSupplier.connection(tenantId);
   }
 
   @Override

ebean-test/src/test/java/org/multitenant/partition/MultiTenantPartitionTest.java

@@ -5,6 +5,9 @@
 import io.ebean.DatabaseBuilder;
 import io.ebean.config.DatabaseConfig;
 import io.ebean.config.TenantMode;
+import io.ebean.meta.MetaQueryPlan;
+import io.ebean.meta.QueryPlanInit;
+import io.ebean.meta.QueryPlanRequest;
 import io.ebean.test.LoggedSql;
 import io.ebean.xtest.BaseTestCase;
 import org.junit.jupiter.api.AfterAll;
@@ -23,12 +26,13 @@ class MultiTenantPartitionTest extends BaseTestCase {
   static List<MtTenant> tenants() {
     List<MtTenant> tenants = new ArrayList<>();
     for (int i = 0; i < 5; i++) {
-      tenants.add(new MtTenant("ten_"+i, names[i], names[i]+"@foo.com".toLowerCase()));
+      tenants.add(new MtTenant("ten_" + i, names[i], names[i] + "@foo.com".toLowerCase()));
     }
     return tenants;
   }
 
   private static final Database server = init();
+
   static {
     server.saveAll(tenants());
   }
@@ -77,6 +81,41 @@ void start() {
     LoggedSql.stop();
   }
 
+  @Test
+  void queryPlanCapture() throws InterruptedException {
+
+    QueryPlanRequest request = new QueryPlanRequest();
+    request.maxCount(1_000);
+    request.maxTimeMillis(10_000);
+    server.metaInfo().queryPlanCollectNow(request);
+
+    QueryPlanInit init = new QueryPlanInit();
+    init.setAll(true);
+    init.thresholdMicros(1);
+    server.metaInfo().queryPlanInit(init);
+
+    try {
+
+      // run queries again
+      UserContext.set("rob", "ten_1");
+      server.find(MtContent.class).findList();
+
+      UserContext.set("fred", "ten_2");
+      server.find(MtContent.class).setId(2).findOne();
+
+      // obtains db query plans ...
+      List<MetaQueryPlan> plans0 = server.metaInfo().queryPlanCollectNow(request);
+      assertThat(plans0).isNotEmpty();
+
+      assertThat(plans0).extracting(MetaQueryPlan::tenantId).containsExactlyInAnyOrder("ten_1", "ten_2");
+    } finally {
+      // disable capturing
+      init.thresholdMicros(Long.MAX_VALUE);
+      server.metaInfo().queryPlanInit(init);
+    }
+  }
+
+
   @Test
   void deleteById() {
     UserContext.set("fred", "ten_2");