Merge pull request #2112 from ebean-orm/FOCONIS-bug-ebean/sql_update_batch_flush

2110 - Fix and test for executeBatch() on SqlUpdate

ebean-core/src/main/java/io/ebeaninternal/server/persist/BatchControl.java

@@ -149,7 +149,6 @@ public final class BatchControl {
       // execute the request immediately without batching
       return request.executeNow();
     }
-
     if (pstmtHolder.getMaxSize() >= batchSize) {
       flush();
     }
@@ -166,7 +165,6 @@ public final class BatchControl {
    * according to the depth (object graph depth).
    */
   public int executeOrQueue(PersistRequestBean<?> request, boolean batch) throws BatchedSqlException {
-
     if (!batch || (batchFlushOnMixed && !pstmtHolder.isEmpty())) {
       // flush when mixing beans and updateSql
       flush();
@@ -185,7 +183,6 @@ public final class BatchControl {
    * Add the request to the batch and return true if we should flush.
    */
   private boolean addToBatch(PersistRequestBean<?> request) {
-
     Object alreadyInBatch = persistedBeans.put(request.getEntityBean(), DUMMY);
     if (alreadyInBatch != null) {
       // special case where the same bean instance has already been
@@ -220,7 +217,11 @@ public final class BatchControl {
    * Flush any batched PreparedStatements.
    */
   private void flushPstmtHolder() throws BatchedSqlException {
-    pstmtHolder.flush(getGeneratedKeys);
+    pstmtHolder.flush(getGeneratedKeys, false);
+  }
+
+  private void flushPstmtHolder(boolean reset) throws BatchedSqlException {
+    pstmtHolder.flush(getGeneratedKeys, reset);
   }
 
   /**
@@ -237,6 +238,15 @@ public final class BatchControl {
     flushPstmtHolder();
   }
 
+  public void flushOnCommit() throws BatchedSqlException {
+    try {
+      flushBuffer(false);
+    } finally {
+      // ensure PreparedStatements are closed
+      pstmtHolder.clear();
+    }
+  }
+
   /**
    * Flush without resetting the depth info.
    */
@@ -261,8 +271,8 @@ public final class BatchControl {
     persistedBeans.clear();
   }
 
-  private void flushBuffer(boolean resetTop) throws BatchedSqlException {
-    flushInternal(resetTop);
+  private void flushBuffer(boolean reset) throws BatchedSqlException {
+    flushInternal(reset);
     flushQueue(earlyQueue);
     flushQueue(lateQueue);
   }
@@ -275,14 +285,15 @@ public final class BatchControl {
 
   /**
    * execute all the requests currently queued or batched.
+   *
+   * @param reset When true close all batched statements (completely empty)
    */
-  private void flushInternal(boolean resetTop) throws BatchedSqlException {
-
+  private void flushInternal(boolean reset) throws BatchedSqlException {
     try {
       bufferMax = 0;
       if (!pstmtHolder.isEmpty()) {
         // Flush existing pstmts (updateSql or callableSql)
-        flushPstmtHolder();
+        flushPstmtHolder(reset);
       }
       if (isEmpty()) {
         // Nothing in queue to flush
@@ -301,7 +312,7 @@ public final class BatchControl {
         beanHolder.executeNow();
       }
       persistedBeans.clear();
-      if (resetTop) {
+      if (reset) {
         beanHoldMap.clear();
         depthOrder.clear();
       }

ebean-core/src/main/java/io/ebeaninternal/server/persist/BatchedPstmt.java

@@ -37,7 +37,7 @@ public class BatchedPstmt implements SpiProfileTransactionEvent {
   /**
    * The list of BatchPostExecute used to perform post processing.
    */
-  private final ArrayList<BatchPostExecute> list = new ArrayList<>();
+  private final List<BatchPostExecute> list = new ArrayList<>();
 
   private final String sql;
 
@@ -68,6 +68,10 @@ public class BatchedPstmt implements SpiProfileTransactionEvent {
     return list.size();
   }
 
+  public boolean isEmpty() {
+    return list.isEmpty();
+  }
+
   /**
    * Return the sql
    */
@@ -95,9 +99,7 @@ public class BatchedPstmt implements SpiProfileTransactionEvent {
     if (rows.length != list.size()) {
       throw new IllegalStateException("Invalid state on executeBatch, rows:" + rows.length + " != " + list.size());
     }
-    for (BatchPostExecute item : list) {
-      item.postExecute();
-    }
+    postExecute();
     list.clear();
   }
 
@@ -113,7 +115,9 @@ public class BatchedPstmt implements SpiProfileTransactionEvent {
    * Run any post processing including getGeneratedKeys.
    */
   public void executeBatch(boolean getGeneratedKeys) throws SQLException {
-
+    if (list.isEmpty()) {
+      return;
+    }
     timedStart = System.nanoTime();
     profileStart = transaction.profileOffset();
     executeAndCheckRowCounts();
@@ -121,8 +125,8 @@ public class BatchedPstmt implements SpiProfileTransactionEvent {
       getGeneratedKeys();
     }
     postExecute();
-    close();
     addTimingMetrics();
+    list.clear();
     transaction.profileEvent(this);
   }
 
@@ -140,10 +144,15 @@ public class BatchedPstmt implements SpiProfileTransactionEvent {
   /**
    * Close the underlying statement.
    */
-  public void close() throws SQLException {
+  public void close() {
     if (pstmt != null) {
-      pstmt.close();
-      pstmt = null;
+      try {
+        pstmt.close();
+      } catch (SQLException e) {
+        log.warn("Error closing statement", e);
+      } finally {
+        pstmt = null;
+      }
     }
   }
 
@@ -169,7 +178,6 @@ public class BatchedPstmt implements SpiProfileTransactionEvent {
   }
 
   private void getGeneratedKeys() throws SQLException {
-
     int index = 0;
     try (ResultSet rset = pstmt.getGeneratedKeys()) {
       while (rset.next()) {
@@ -190,8 +198,13 @@ public class BatchedPstmt implements SpiProfileTransactionEvent {
   /**
    * Register any inputStreams that should be closed after execution.
    */
-  public void registerInputStreams(List<InputStream> inputStreams) {
-    this.inputStreams = inputStreams;
+  public void registerInputStreams(List<InputStream> streams) {
+    if (streams != null) {
+      if (this.inputStreams == null) {
+        this.inputStreams = new ArrayList<>();
+      }
+      this.inputStreams.addAll(streams);
+    }
   }
 
   private void closeInputStreams() {
@@ -203,6 +216,7 @@ public class BatchedPstmt implements SpiProfileTransactionEvent {
           log.warn("Error closing inputStream ", e);
         }
       }
+      inputStreams = null;
     }
   }
 }

ebean-core/src/main/java/io/ebeaninternal/server/persist/BatchedPstmtHolder.java

@@ -1,12 +1,11 @@
 package io.ebeaninternal.server.persist;
 
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
 import javax.persistence.PersistenceException;
 import java.sql.PreparedStatement;
 import java.sql.SQLException;
+import java.util.Collection;
 import java.util.LinkedHashMap;
+import java.util.Map;
 
 /**
  * Used to hold BatchedPstmt objects for batch based execution.
@@ -18,13 +17,11 @@ import java.util.LinkedHashMap;
  */
 public class BatchedPstmtHolder {
 
-  private static final Logger logger = LoggerFactory.getLogger(BatchedPstmtHolder.class);
-
   /**
    * A Map of the statements using a String key. This is used so that the same
    * Statement,Prepared,Callable is reused.
    */
-  private final LinkedHashMap<String, BatchedPstmt> stmtMap = new LinkedHashMap<>();
+  private Map<String, BatchedPstmt> stmtMap = new LinkedHashMap<>();
 
   /**
    * The Max size across all the BatchedPstmt.
@@ -47,13 +44,10 @@ public class BatchedPstmtHolder {
    * Return the BatchedPstmt that holds the batched statement.
    */
   public BatchedPstmt getBatchedPstmt(String stmtKey) {
-
     BatchedPstmt bs = stmtMap.get(stmtKey);
     if (bs == null) {
-      // the PreparedStatement has need been created
       return null;
     }
-
     // maintain a max batch size for any given batched stmt.
     // Used to determine when to flush.
     int bsSize = bs.size();
@@ -63,14 +57,19 @@ public class BatchedPstmtHolder {
     return bs;
   }
 
+  /**
+   * Return the size of the biggest batched statement.
+   * Used to determine when to flush the batch.
+   */
+  int getMaxSize() {
+    return maxSize;
+  }
+
   /**
    * Add a new PreparedStatement wrapped in the BatchStatement object.
    */
   public void addStmt(BatchedPstmt bs, BatchPostExecute postExecute) {
-    // add the batch post execute to the statement for POST processing
     bs.add(postExecute);
-
-    // cache so that getStmt() can find it for additional beans/rows
     stmtMap.put(bs.getSql(), bs);
   }
 
@@ -78,15 +77,22 @@ public class BatchedPstmtHolder {
    * Return true if the batch has no statements to execute.
    */
   public boolean isEmpty() {
-    return stmtMap.isEmpty();
+    if (stmtMap.isEmpty()) {
+      return true;
+    }
+    for (BatchedPstmt bs : stmtMap.values()) {
+      if (!bs.isEmpty()) {
+        return false;
+      }
+    }
+    return true;
   }
 
   /**
    * Execute one of the batched statements returning the row counts.
    */
   public int[] execute(String key, boolean getGeneratedKeys) throws SQLException {
-
-    BatchedPstmt batchedPstmt = stmtMap.remove(key);
+    BatchedPstmt batchedPstmt = stmtMap.get(key);
     if (batchedPstmt == null) {
       throw new PersistenceException("No batched statement found for key " + key);
     }
@@ -96,69 +102,56 @@ public class BatchedPstmtHolder {
 
   /**
    * Execute all batched PreparedStatements.
-   *
-   * @param getGeneratedKeys if true try to get generated keys for inserts
    */
-  public void flush(boolean getGeneratedKeys) throws BatchedSqlException {
-
-    SQLException firstError = null;
-    String errorSql = null;
-
-    // flag set if something fails. Will not execute
-    // but still need to close PreparedStatements.
-    boolean isError = false;
-
+  public void flush(boolean getGeneratedKeys, boolean reset) throws BatchedSqlException {
     // if there are Listeners/Controllers that interact with the database,
     // the flush may get called recursively in executeBatch/postExecute.
-    // which leads that we process stmtMap.values() twice in the loop.
-    // So we copy the values, that we want to flush and clear it immediately.
-    BatchedPstmt[] values = stmtMap.values().toArray(new BatchedPstmt[0]);
-    clear();
+    // which means this needs to process a copy of stmtMap, create a new stmtMap and loadBack after
+    final Map<String, BatchedPstmt> copyMap = stmtMap;
+    final Collection<BatchedPstmt> copy = copyMap.values();
+    this.stmtMap = new LinkedHashMap<>();
+    this.maxSize = 0;
+    try {
+      executeAll(copy, getGeneratedKeys);
+      if (reset) {
+        closeStatements(copy);
+      } else {
+        loadBack(copyMap);
+      }
+    } catch (BatchedSqlException e) {
+      closeStatements(copy);
+      throw e;
+    }
+  }
 
+  private void loadBack(Map<String, BatchedPstmt> copyMap) {
+    if (stmtMap.isEmpty()) {
+      // just restore, was not modified during flush by Listeners/Controllers
+      stmtMap = copyMap;
+    } else {
+      closeStatements(copyMap.values());
+    }
+  }
+
+  private void executeAll(Collection<BatchedPstmt> values, boolean getGeneratedKeys) throws BatchedSqlException {
     for (BatchedPstmt bs : values) {
       try {
-        if (!isError) {
-          bs.executeBatch(getGeneratedKeys);
-        }
+        bs.executeBatch(getGeneratedKeys);
       } catch (SQLException ex) {
-        SQLException next = ex.getNextException();
-        while (next != null) {
-          logger.trace("Next Exception during batch execution", next);
-          next = next.getNextException();
-        }
-
-        firstError = ex;
-        errorSql = bs.getSql();
-        isError = true;
-
-      } finally {
-        try {
-          bs.close();
-        } catch (SQLException ex) {
-          logger.error("Error closing batched PreparedStatement", ex);
-        }
+        throw new BatchedSqlException("Error when batch flush on sql: " + bs.getSql(), ex);
       }
     }
-
-    if (firstError != null) {
-      String msg = "Error when batch flush on sql: " + errorSql;
-      throw new BatchedSqlException(msg, firstError);
-    }
   }
 
   public void clear() {
+    closeStatements(stmtMap.values());
     stmtMap.clear();
     maxSize = 0;
   }
 
-  /**
-   * Return the size of the biggest batched statement.
-   * <p>
-   * Used to determine when to flush the batch.
-   * </p>
-   */
-  int getMaxSize() {
-    return maxSize;
+  private void closeStatements(Collection<BatchedPstmt> batchedStatements) {
+    for (BatchedPstmt bs: batchedStatements) {
+      bs.close();
+    }
   }
-
 }

ebean-core/src/main/java/io/ebeaninternal/server/persist/PstmtFactory.java

@@ -44,10 +44,12 @@ class PstmtFactory {
    * Return a prepared statement taking into account batch requirements.
    */
   PreparedStatement getPstmtBatch(SpiTransaction t, String sql, BatchPostExecute batchExe) throws SQLException {
-
     BatchedPstmtHolder batch = t.getBatchControl().getPstmtHolder();
     BatchedPstmt existingStmt = batch.getBatchedPstmt(sql);
     if (existingStmt != null) {
+      if (existingStmt.isEmpty() && t.isLogSql()) {
+        t.logSql(TrimLogSql.trim(sql));
+      }
       return existingStmt.getStatement(batchExe);
     }
 
@@ -57,7 +59,6 @@ class PstmtFactory {
 
     Connection conn = t.getInternalConnection();
     PreparedStatement stmt = conn.prepareStatement(sql);
-
     BatchedPstmt bs = new BatchedPstmt(stmt, false, sql, t);
     batch.addStmt(bs, batchExe);
     return stmt;

ebean-core/src/main/java/io/ebeaninternal/server/persist/dml/DmlHandler.java

@@ -24,8 +24,8 @@ public abstract class DmlHandler implements PersistHandler, BindableRequest {
   private static final Logger logger = LoggerFactory.getLogger(DmlHandler.class);
 
   private static final int[] GENERATED_KEY_COLUMNS = new int[]{1};
-  private static final int BATCHED_FIRST = 1;
-  private static final int BATCHED = 2;
+  private static final short BATCHED_FIRST = 1;
+  private static final short BATCHED = 2;
 
   /**
    * The originating request.
@@ -261,17 +261,15 @@ public abstract class DmlHandler implements PersistHandler, BindableRequest {
    * Return a prepared statement taking into account batch requirements.
    */
   PreparedStatement getPstmtBatch(SpiTransaction t, String sql, PersistRequestBean<?> request, boolean genKeys) throws SQLException {
-
     BatchedPstmtHolder batch = t.getBatchControl().getPstmtHolder();
     batchedPstmt = batch.getBatchedPstmt(sql);
     if (batchedPstmt != null) {
-      batchedStatus = BATCHED;
+      batchedStatus = batchedPstmt.isEmpty() ? BATCHED_FIRST : BATCHED;
       return batchedPstmt.getStatement(request);
     }
 
     batchedStatus = BATCHED_FIRST;
     PreparedStatement stmt = getPstmt(t, sql, genKeys);
-
     batchedPstmt = new BatchedPstmt(stmt, genKeys, sql, t);
     batch.addStmt(batchedPstmt, request);
     return stmt;

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

@@ -131,7 +131,7 @@ class JdbcTransaction implements SpiTransaction, TxnProfileEventCodes {
 
   private Boolean batchFlushOnMixed;
 
-  private String logPrefix;
+  private final String logPrefix;
 
   private Object tenantId;
 
@@ -649,7 +649,7 @@ class JdbcTransaction implements SpiTransaction, TxnProfileEventCodes {
   private void batchFlush() {
     if (batchControl != null) {
       try {
-        batchControl.flush();
+        batchControl.flushOnCommit();
       } catch (BatchedSqlException e) {
         throw translate(e.getMessage(), e.getCause());
       }
@@ -686,16 +686,20 @@ class JdbcTransaction implements SpiTransaction, TxnProfileEventCodes {
 
   @Override
   public void flushBatchOnRollback() {
-    if (batchControl != null) {
-      if (logger.isTraceEnabled()) {
-        logger.trace("... flushBatchOnRollback");
-      }
-      batchControl.clear();
-    }
+    internalBatchClear();
     // restore the previous batch mode
     batchMode = oldBatchMode;
   }
 
+  /**
+   * Ensure batched PreparedStatements are closed on rollback.
+   */
+  private void internalBatchClear() {
+    if (batchControl != null) {
+      batchControl.clear();
+    }
+  }
+
   @Override
   public boolean checkBatchEscalationOnCascade(PersistRequestBean<?> request) {
 
@@ -1119,6 +1123,7 @@ class JdbcTransaction implements SpiTransaction, TxnProfileEventCodes {
    * Perform the jdbc rollback and fire any registered callbacks.
    */
   private void doRollback(Throwable cause) {
+    internalBatchClear();
     firePreRollback();
     try {
       performRollback();

ebean-core/src/test/java/org/tests/cascade/TestMultiCascadeBatch.java

@@ -1,6 +1,7 @@
 package org.tests.cascade;
 
 import io.ebean.BaseTestCase;
+import io.ebean.DB;
 import io.ebean.Ebean;
 import io.ebean.Transaction;
 import org.ebeantest.LoggedSqlCollector;
@@ -49,6 +50,12 @@ public class TestMultiCascadeBatch extends BaseTestCase {
 
     final List<String> sql = LoggedSqlCollector.stop();
 
+    final List<Site> list = DB.find(Site.class).where()
+      .idIn(grandparent.getId(), parent.getId(), child.getId())
+      .findList();
+
+    assertThat(list).hasSize(3);
+
     assertThat(sql).hasSize(5);
     assertSql(sql.get(0)).contains("insert into site (id, name");
     assertSql(sql.get(1)).contains("insert into site (id, name");

ebean-core/src/test/java/org/tests/update/TestSqlUpdateBatch.java

@@ -0,0 +1,67 @@
+package org.tests.update;
+
+import io.ebean.BaseTestCase;
+import io.ebean.DB;
+import io.ebean.SqlUpdate;
+import io.ebean.Transaction;
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+
+/**
+ * Testclass that ensures the correct behaviour of SqlUpdate in combination with batching.
+ *
+ * @author Jonas Pöhler, FOCONIS AG
+ */
+public class TestSqlUpdateBatch extends BaseTestCase {
+
+  /**
+   * This test ensures the correct behaviour of update batching in the {@code BatchedPstmtHolder} after a flush. Two batch updates
+   * are executed with each a batchsize of 21, which is the default limit for automatic flushing on {@code addBatch()}. After that
+   * the statements are each executed with {@code executeBatch()}. The desired behaviour is, that each execution succeeds although
+   * technically the batch was already flushed.
+   */
+  @Test
+  public void testTwoParallelBatches() {
+    try (Transaction txn = DB.beginTransaction()) {
+      // Dummy updates, that effectively do nothing, but ebean doesn't need to know this.
+      final SqlUpdate update = DB.sqlUpdate("update uuone set name = ? where 0=1");
+      final SqlUpdate delete = DB.sqlUpdate("delete from uuone where ?=-1");
+      // txn.setBatchSize(40);
+
+      for (int i = 0; i <= 20; i++) {
+        update
+          .setParameter(1, String.valueOf(i))
+          .addBatch();
+        delete
+          .setParameter(1, String.valueOf(i))
+          .addBatch();
+      }
+
+      delete.executeBatch();
+      update.executeBatch();
+      txn.commit();
+    }
+  }
+
+  /**
+   * This test checks that for a batch update a correct array with update counts is returned. If a update with 40 entries is
+   * executed, it is expected, that {@code executeBatch()} returns a array with 40 elements each containing the update count for
+   * the given parameters.
+   */
+  @Test
+  public void testBatchReturnArrayLength() {
+    try (Transaction txn = DB.beginTransaction()) {
+      txn.setBatchSize(100); // something bigger than 40
+      // Dummy update, that effectively does nothing, but ebean doesn't need to know this.
+      final SqlUpdate update = DB.sqlUpdate("update uuone set name = ? where 0=1");
+      for (int i = 0; i < 40; i++) {
+        update
+          .setParameter(1, String.valueOf(i))
+          .addBatch();
+      }
+      assertEquals(40, update.executeBatch().length);
+    }
+  }
+
+}