From 5ddcb64d17f8a44ca4c3e5e4f6fc4d040bb40a59 Mon Sep 17 00:00:00 2001
From: Radim Vansa <rvansa@azul.com>
Date: Wed, 29 Mar 2023 14:51:23 +0200
Subject: [PATCH] Close connections during CRaC snapshotting

Signed-off-by: Radim Vansa <rvansa@azul.com>
---
 agroal-pool/pom.xml                           |  5 ++
 .../java/io/agroal/pool/ConnectionPool.java   | 21 +++++-
 .../io/agroal/test/basic/CheckpointTests.java | 73 +++++++++++++++++++
 pom.xml                                       |  1 +
 4 files changed, 98 insertions(+), 2 deletions(-)
 create mode 100644 agroal-test/src/test/java/io/agroal/test/basic/CheckpointTests.java

diff --git a/agroal-pool/pom.xml b/agroal-pool/pom.xml
index adccfd17..b1a85748 100644
--- a/agroal-pool/pom.xml
+++ b/agroal-pool/pom.xml
@@ -15,6 +15,11 @@
             <artifactId>agroal-api</artifactId>
             <version>${forked.version}</version>
         </dependency>
+        <dependency>
+            <groupId>org.crac</groupId>
+            <artifactId>crac</artifactId>
+            <version>${version.org.crac}</version>
+        </dependency>
     </dependencies>
     <build>
         <plugins>
diff --git a/agroal-pool/src/main/java/io/agroal/pool/ConnectionPool.java b/agroal-pool/src/main/java/io/agroal/pool/ConnectionPool.java
index 62e5348c..c875b9ed 100644
--- a/agroal-pool/src/main/java/io/agroal/pool/ConnectionPool.java
+++ b/agroal-pool/src/main/java/io/agroal/pool/ConnectionPool.java
@@ -31,6 +31,10 @@
 import java.util.concurrent.atomic.LongAdder;
 import java.util.function.Function;
 
+import org.crac.Context;
+import org.crac.Core;
+import org.crac.Resource;
+
 import static io.agroal.api.AgroalDataSource.FlushMode.GRACEFUL;
 import static io.agroal.api.AgroalDataSource.FlushMode.LEAK;
 import static io.agroal.api.configuration.AgroalConnectionPoolConfiguration.MultipleAcquisitionAction.OFF;
@@ -74,7 +78,7 @@
 /**
  * @author <a href="lbarreiro@redhat.com">Luis Barreiro</a>
  */
-public final class ConnectionPool implements Pool {
+public final class ConnectionPool implements Pool, Resource {
 
     private static final AtomicInteger HOUSEKEEP_COUNT = new AtomicInteger();
 
@@ -85,7 +89,7 @@ public final class ConnectionPool implements Pool {
 
     private final AgroalSynchronizer synchronizer;
     private final ConnectionFactory connectionFactory;
-    private final PriorityScheduledExecutor housekeepingExecutor;
+    private PriorityScheduledExecutor housekeepingExecutor;
     private final TransactionIntegration transactionIntegration;
 
     private final boolean borrowValidationEnabled;
@@ -118,6 +122,8 @@ public ConnectionPool(AgroalConnectionPoolConfiguration configuration, AgroalDat
         leakEnabled = !configuration.leakTimeout().isZero();
         validationEnabled = !configuration.validationTimeout().isZero();
         reapEnabled = !configuration.reapTimeout().isZero();
+
+        Core.getGlobalContext().register(this);
     }
 
     public void init() {
@@ -556,6 +562,17 @@ public boolean isHealthy(boolean newConnection) throws SQLException {
         return performValidation( healthHandler, CHECKED_IN );
     }
 
+    @Override
+    public void beforeCheckpoint(Context<? extends Resource> context) throws Exception {
+        close();
+    }
+
+    @Override
+    public void afterRestore(Context<? extends Resource> context) throws Exception {
+        housekeepingExecutor = new PriorityScheduledExecutor( 1, "agroal-" + HOUSEKEEP_COUNT.incrementAndGet(), listeners );
+        init();
+    }
+
     // --- create //
 
     private final class CreateConnectionTask implements Callable<ConnectionHandler> {
diff --git a/agroal-test/src/test/java/io/agroal/test/basic/CheckpointTests.java b/agroal-test/src/test/java/io/agroal/test/basic/CheckpointTests.java
new file mode 100644
index 00000000..fcbb7726
--- /dev/null
+++ b/agroal-test/src/test/java/io/agroal/test/basic/CheckpointTests.java
@@ -0,0 +1,73 @@
+package io.agroal.test.basic;
+
+import io.agroal.api.AgroalDataSource;
+import io.agroal.api.configuration.supplier.AgroalDataSourceConfigurationSupplier;
+import io.agroal.pool.DataSource;
+import io.agroal.pool.Pool;
+import io.agroal.test.MockConnection;
+import org.crac.Resource;
+import org.junit.jupiter.api.*;
+
+import java.lang.reflect.Field;
+import java.sql.Connection;
+
+import static io.agroal.test.AgroalTestGroup.FUNCTIONAL;
+import static io.agroal.test.MockDriver.deregisterMockDriver;
+import static io.agroal.test.MockDriver.registerMockDriver;
+import static org.junit.jupiter.api.Assertions.*;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+@Tag( FUNCTIONAL )
+public class CheckpointTests {
+    @BeforeAll
+    static void setupMockDriver() {
+        registerMockDriver( FakeConnection.class );
+    }
+
+    @AfterAll
+    static void teardown() {
+        deregisterMockDriver();
+    }
+
+    // --- //
+
+    @Test
+    @DisplayName( "ConnectionPool C/R" )
+    void checkpointRestoreTest() throws Exception {
+        try ( AgroalDataSource dataSource = AgroalDataSource.from( new AgroalDataSourceConfigurationSupplier().connectionPoolConfiguration(cp -> cp.maxSize( 1 ) ) ) ) {
+            Connection connection = dataSource.getConnection();
+
+            assertFalse( connection.isClosed(), "Expected open connection, but it's closed" );
+
+            Field cpField = DataSource.class.getDeclaredField("connectionPool");
+            cpField.setAccessible(true);
+            Pool pool = (Pool) cpField.get(dataSource);
+            assertInstanceOf(Resource.class, pool);
+            Resource poolResource = (Resource) pool;
+
+            poolResource.beforeCheckpoint(null);
+
+            assertTrue( connection.isClosed(), "Expected closed connection, but it's open" );
+
+            poolResource.afterRestore(null);
+
+            Connection another = dataSource.getConnection();
+            assertFalse( another.isClosed(), "Expected open connection, but it's closed" );
+        }
+    }
+
+    public static class FakeConnection implements MockConnection {
+        private boolean closed;
+
+        @Override
+        public void close() {
+            closed = true;
+        }
+
+        @Override
+        public boolean isClosed() {
+            return closed;
+        }
+    }
+
+}
diff --git a/pom.xml b/pom.xml
index d8f077c8..fa8d74af 100644
--- a/pom.xml
+++ b/pom.xml
@@ -83,6 +83,7 @@
         <version.eu.somatik.serviceloader-maven-plugin>1.4.0</version.eu.somatik.serviceloader-maven-plugin>
         <version.org.apache.maven.plugins.maven-surefire-plugin>3.3.1</version.org.apache.maven.plugins.maven-surefire-plugin>
         <version.org.netbeans.tools.sigtest-maven-plugin>1.6</version.org.netbeans.tools.sigtest-maven-plugin>
+        <version.org.crac>1.5.0</version.org.crac>
     </properties>
     <build>
         <plugins>