framework/db: keep the caller's transaction when persist() fails

framework/db/src/main/java/com/cloud/utils/db/GenericDaoBase.java

@@ -33,6 +33,7 @@
 import java.sql.ResultSet;
 import java.sql.ResultSetMetaData;
 import java.sql.SQLException;
+import java.sql.Savepoint;
 import java.sql.Statement;
 import java.sql.Time;
 import java.sql.Timestamp;
@@ -1632,10 +1633,18 @@ public T persist(final T entity) {
 
         ID id = null;
         final TransactionLegacy txn = TransactionLegacy.currentTxn();
+        final boolean nested = txn.dbTxnStarted();
+        Savepoint savepoint = null;
+        boolean committed = false;
         PreparedStatement pstmt = null;
         String sql = null;
         try {
-            txn.start();
+            if (nested) {
+                // Inside the caller's transaction, so a failure must only undo this insert
+                savepoint = txn.setSavepoint();
+            } else {
+                txn.start();
+            }
             for (final Pair<String, Attribute[]> pair : _insertSqls) {
                 sql = pair.first();
                 final Attribute[] attrs = pair.second();
@@ -1681,7 +1690,12 @@ public T persist(final T entity) {
 
                 insertElementCollection(entity, _idAttributes.get(_table)[0], id, ecAttributes);
             }
-            txn.commit();
+            if (nested) {
+                txn.releaseSavepoint(savepoint);
+            } else {
+                txn.commit();
+            }
+            committed = true;
         } catch (final SQLException e) {
             logger.error("DB Exception on: " + pstmt, e);
             handleEntityExistsException(e);
@@ -1690,6 +1704,14 @@ public T persist(final T entity) {
             throw new CloudRuntimeException("Problem with getting the ec attribute ", e);
         } catch (IllegalAccessException e) {
             throw new CloudRuntimeException("Problem with getting the ec attribute ", e);
+        } finally {
+            if (!committed) {
+                if (!nested) {
+                    txn.rollback();
+                } else if (savepoint != null) {
+                    txn.rollback(savepoint);
+                }
+            }
         }
 
         return _idField != null ? findByIdIncludingRemoved(id) : null;

framework/db/src/main/java/com/cloud/utils/db/TransactionLegacy.java

@@ -898,31 +898,27 @@ public void rollback() {
 
     public Savepoint setSavepoint() throws SQLException {
         _txn = true;
-        StackElement st = new StackElement(START_TXN, null);
-        _stack.push(st);
         final Connection conn = getConnection();
         final Savepoint sp = conn.setSavepoint();
-        st.ref = sp;
+        _stack.push(new StackElement(START_TXN, sp));
 
         return sp;
     }
 
     public Savepoint setSavepoint(final String name) throws SQLException {
         _txn = true;
-        StackElement st = new StackElement(START_TXN, null);
-        _stack.push(st);
         final Connection conn = getConnection();
         final Savepoint sp = conn.setSavepoint(name);
-        st.ref = sp;
+        _stack.push(new StackElement(START_TXN, sp));
 
         return sp;
     }
 
     public void releaseSavepoint(final Savepoint sp) throws SQLException {
-        removeTxn(sp);
         if (_conn != null) {
             _conn.releaseSavepoint(sp);
         }
+        removeTxn(sp);
 
         if (!hasTxnInStack()) {
             _txn = false;

framework/db/src/test/java/com/cloud/utils/GenericDaoBasePersistTransactionTest.java

@@ -0,0 +1,216 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// the License.  You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+package com.cloud.utils;
+
+import java.sql.Connection;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Savepoint;
+import java.sql.Statement;
+
+import javax.persistence.Column;
+import javax.persistence.Entity;
+import javax.persistence.EntityExistsException;
+import javax.persistence.GeneratedValue;
+import javax.persistence.GenerationType;
+import javax.persistence.Id;
+import javax.persistence.Table;
+
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.ArgumentMatchers;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.junit.MockitoJUnitRunner;
+
+import com.cloud.utils.db.GenericDaoBase;
+import com.cloud.utils.db.TransactionLegacy;
+import com.cloud.utils.exception.CloudRuntimeException;
+
+// Kept outside com.cloud.utils.db because framework/db excludes that package's tests from the build.
+@RunWith(MockitoJUnitRunner.class)
+public class GenericDaoBasePersistTransactionTest {
+
+    private static final String TXN_NAME = "GenericDaoBasePersistTransactionTest";
+
+    @Mock
+    private Connection conn;
+    @Mock
+    private PreparedStatement insert;
+    @Mock
+    private PreparedStatement select;
+    @Mock
+    private ResultSet noRows;
+    @Mock
+    private Savepoint savepoint;
+
+    private final PersistTestDao dao = new PersistTestDao();
+    private TransactionLegacy txn;
+
+    @Before
+    public void setUp() {
+        txn = TransactionLegacy.open(TXN_NAME);
+        txn.transitToUserManagedConnection(conn);
+    }
+
+    @After
+    public void tearDown() {
+        // A failing test can leave the thread's transaction open, which would leak into the next test
+        while (txn.dbTxnStarted()) {
+            txn.rollback();
+        }
+        txn.close(TXN_NAME);
+        txn.transitToAutoManagedConnection(TransactionLegacy.CLOUD_DB);
+    }
+
+    private void insertSucceeds() throws SQLException {
+        Mockito.when(conn.prepareStatement(ArgumentMatchers.anyString(), ArgumentMatchers.eq(Statement.RETURN_GENERATED_KEYS))).thenReturn(insert);
+        Mockito.when(conn.prepareStatement(ArgumentMatchers.anyString())).thenReturn(select);
+        Mockito.when(select.executeQuery()).thenReturn(noRows);
+    }
+
+    private void insertFailsWithDuplicateEntry() throws SQLException {
+        Mockito.when(conn.prepareStatement(ArgumentMatchers.anyString(), ArgumentMatchers.eq(Statement.RETURN_GENERATED_KEYS))).thenReturn(insert);
+        Mockito.when(insert.executeUpdate()).thenThrow(new SQLException("Duplicate entry", "23000", 1062));
+    }
+
+    @Test
+    public void failedPersistInsideCallerTransactionOnlyRollsBackToItsSavepoint() throws SQLException {
+        Mockito.when(conn.setSavepoint()).thenReturn(savepoint);
+        insertFailsWithDuplicateEntry();
+        txn.start();
+
+        Assert.assertThrows(EntityExistsException.class, () -> dao.persist(new PersistTestVO()));
+
+        Mockito.verify(conn).rollback(savepoint);
+        Mockito.verify(conn, Mockito.never()).rollback();
+        Assert.assertTrue(txn.dbTxnStarted());
+        Assert.assertTrue(txn.commit());
+        Mockito.verify(conn).commit();
+    }
+
+    @Test
+    public void failedPersistOutsideATransactionRollsBackItsOwnTransaction() throws SQLException {
+        insertFailsWithDuplicateEntry();
+
+        Assert.assertThrows(EntityExistsException.class, () -> dao.persist(new PersistTestVO()));
+
+        Mockito.verify(conn).rollback();
+        Mockito.verify(conn, Mockito.never()).setSavepoint();
+        Assert.assertFalse(txn.dbTxnStarted());
+    }
+
+    @Test
+    public void persistInsideCallerTransactionReleasesItsSavepointAndLeavesTheCommitToTheCaller() throws SQLException {
+        Mockito.when(conn.setSavepoint()).thenReturn(savepoint);
+        insertSucceeds();
+        txn.start();
+
+        dao.persist(new PersistTestVO());
+
+        Mockito.verify(conn).releaseSavepoint(savepoint);
+        Mockito.verify(conn, Mockito.never()).commit();
+        Assert.assertTrue(txn.commit());
+        Mockito.verify(conn).commit();
+    }
+
+    @Test
+    public void persistOutsideATransactionCommitsWithoutASavepoint() throws SQLException {
+        insertSucceeds();
+
+        dao.persist(new PersistTestVO());
+
+        Mockito.verify(conn).commit();
+        Mockito.verify(conn, Mockito.never()).setSavepoint();
+        Assert.assertFalse(txn.dbTxnStarted());
+    }
+
+    @Test
+    public void persistInsideCallerTransactionLeavesItIntactWhenTheSavepointCannotBeTaken() throws SQLException {
+        Mockito.when(conn.setSavepoint()).thenThrow(new SQLException("Connection lost"));
+        txn.start();
+
+        Assert.assertThrows(CloudRuntimeException.class, () -> dao.persist(new PersistTestVO()));
+
+        Mockito.verify(conn, Mockito.never()).rollback();
+        Assert.assertTrue(txn.dbTxnStarted());
+        Assert.assertTrue(txn.commit());
+        Mockito.verify(conn).commit();
+    }
+
+    @Test
+    public void failedSetSavepointLeavesTheCallerTransactionIntact() throws SQLException {
+        Mockito.when(conn.setSavepoint()).thenThrow(new SQLException("Connection lost"));
+        txn.start();
+
+        Assert.assertThrows(SQLException.class, () -> txn.setSavepoint());
+
+        Assert.assertTrue(txn.commit());
+        Mockito.verify(conn).commit();
+    }
+
+    @Test
+    public void failedReleaseSavepointLeavesItAvailableForRollback() throws SQLException {
+        Mockito.when(conn.setSavepoint()).thenReturn(savepoint);
+        Mockito.doThrow(new SQLException("Connection lost")).when(conn).releaseSavepoint(savepoint);
+        txn.start();
+        final Savepoint sp = txn.setSavepoint();
+
+        Assert.assertThrows(SQLException.class, () -> txn.releaseSavepoint(sp));
+        txn.rollback(sp);
+
+        Mockito.verify(conn).rollback(savepoint);
+        Assert.assertTrue(txn.commit());
+        Mockito.verify(conn).commit();
+    }
+
+    @Test
+    public void namedSavepointRollsBackOnlyItsOwnLevel() throws SQLException {
+        Mockito.when(conn.setSavepoint("persist")).thenReturn(savepoint);
+        txn.start();
+        final Savepoint sp = txn.setSavepoint("persist");
+
+        txn.rollback(sp);
+
+        Mockito.verify(conn).rollback(savepoint);
+        Mockito.verify(conn, Mockito.never()).rollback();
+        Assert.assertTrue(txn.commit());
+        Mockito.verify(conn).commit();
+    }
+
+    @Entity
+    @Table(name = "persist_test")
+    public static class PersistTestVO {
+        @Id
+        @GeneratedValue(strategy = GenerationType.IDENTITY)
+        @Column(name = "id")
+        private long id;
+
+        @Column(name = "name")
+        private String name;
+
+        public PersistTestVO() {
+        }
+    }
+
+    public static class PersistTestDao extends GenericDaoBase<PersistTestVO, Long> {
+    }
+}