"merged" #2126 Added Support For beginRequest and endRequest (new ver…

…sion)
brettwooldridge · Sep 23, 2024 · d544180 · d544180
1 parent 8053e39
commit d544180
Show file tree

Hide file tree

Showing 3 changed files with 82 additions and 0 deletions.
diff --git a/src/main/java/com/zaxxer/hikari/pool/HikariPool.java b/src/main/java/com/zaxxer/hikari/pool/HikariPool.java
@@ -64,6 +64,7 @@ public final class HikariPool extends PoolBase implements HikariPoolMXBean, IBag
 
    private final long aliveBypassWindowMs = Long.getLong("com.zaxxer.hikari.aliveBypassWindowMs", MILLISECONDS.toMillis(500));
    private final long housekeepingPeriodMs = Long.getLong("com.zaxxer.hikari.housekeeping.periodMs", SECONDS.toMillis(30));
+   private final boolean isRequestBoundariesEnabled = Boolean.getBoolean("com.zaxxer.hikari.enableRequestBoundaries");
 
    private static final String EVICTED_CONNECTION_MESSAGE = "(connection was evicted)";
    private static final String DEAD_CONNECTION_MESSAGE = "(connection is dead)";
@@ -171,6 +172,13 @@ public Connection getConnection(final long hardTimeout) throws SQLException
             }
             else {
                metricsTracker.recordBorrowStats(poolEntry, startTime);
+               if (isRequestBoundariesEnabled) {
+                  try {
+                     poolEntry.connection.beginRequest();
+                  } catch (SQLException e) {
+                     logger.warn("beginRequest Failed for: {}, ({})", poolEntry.connection, e.getMessage());
+                  }
+               }
                return poolEntry.createProxyConnection(leakTaskFactory.schedule(poolEntry));
             }
          } while (timeout > 0L);
@@ -420,6 +428,13 @@ void recycle(final PoolEntry poolEntry)
       if (poolEntry.isMarkedEvicted()) {
          closeConnection(poolEntry, EVICTED_CONNECTION_MESSAGE);
       } else {
+         if (isRequestBoundariesEnabled) {
+            try {
+               poolEntry.connection.endRequest();
+            } catch (SQLException e) {
+               logger.warn("endRequest Failed for: {},({})", poolEntry.connection, e.getMessage());
+            }
+         }
          connectionBag.requite(poolEntry);
       }
    }

diff --git a/src/test/java/com/zaxxer/hikari/mocks/StubConnection.java b/src/test/java/com/zaxxer/hikari/mocks/StubConnection.java
@@ -57,6 +57,8 @@ public class StubConnection extends StubBaseConnection
    private String catalog;
    private String schema;
    private long waitTimeout;
+   public boolean beginRequestCalled = false;
+   public boolean endRequestCalled = false;
 
    private static ScheduledExecutorService connectionWaitTimeout = new ScheduledThreadPoolExecutor(1);
    private ScheduledFuture<?> waitTimeoutTask;
@@ -555,4 +557,13 @@ public int getNetworkTimeout() throws SQLException
       return 0;
    }
 
+   @Override
+   public void beginRequest() {
+      beginRequestCalled = true;
+   }
+
+   @Override
+   public void endRequest() {
+      endRequestCalled = true;
+   }
 }
diff --git a/src/test/java/com/zaxxer/hikari/pool/RequestBoundariesTest.java b/src/test/java/com/zaxxer/hikari/pool/RequestBoundariesTest.java
@@ -0,0 +1,56 @@
+package com.zaxxer.hikari.pool;
+
+import com.zaxxer.hikari.HikariConfig;
+import com.zaxxer.hikari.HikariDataSource;
+import com.zaxxer.hikari.mocks.StubConnection;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.sql.Connection;
+
+import static com.zaxxer.hikari.pool.TestElf.getPool;
+import static com.zaxxer.hikari.pool.TestElf.newHikariConfig;
+
+public class RequestBoundariesTest {
+
+   private static final HikariConfig config;
+   static {
+      config = newHikariConfig();
+      config.setMinimumIdle(0);
+      config.setMaximumPoolSize(10);
+      config.setInitializationFailTimeout(Long.MAX_VALUE);
+      config.setConnectionTestQuery("VALUES 1");
+      config.setDataSourceClassName("com.zaxxer.hikari.mocks.StubDataSource");
+   }
+
+   private HikariPool getHikariPool(boolean enableRequestBoundaries) {
+      System.setProperty("com.zaxxer.hikari.enableRequestBoundaries", String.valueOf(enableRequestBoundaries));
+      HikariDataSource ds = new HikariDataSource(config);
+      HikariPool pool = getPool(ds);
+      return pool;
+   }
+
+   @Test
+   public void requestBoundaryEnabledTest() throws Exception {
+      HikariPool pool = getHikariPool(true);
+      Connection conn = pool.getConnection();
+      StubConnection stubConnection = conn.unwrap(StubConnection.class);
+      Assert.assertTrue("Begin request called", stubConnection.beginRequestCalled);
+      Assert.assertFalse("End request called", stubConnection.endRequestCalled);
+      conn.close();
+      Assert.assertTrue("Begin request called", stubConnection.beginRequestCalled);
+      Assert.assertTrue("End request called", stubConnection.endRequestCalled);
+   }
+
+   @Test
+   public void requestBoundaryDisabledTest() throws Exception {
+      HikariPool pool = getHikariPool(false);
+      Connection conn = pool.getConnection();
+      StubConnection stubConnection = conn.unwrap(StubConnection.class);
+      Assert.assertFalse("Begin request called", stubConnection.beginRequestCalled);
+      Assert.assertFalse("End request called", stubConnection.endRequestCalled);
+      conn.close();
+      Assert.assertFalse("Begin request called", stubConnection.beginRequestCalled);
+      Assert.assertFalse("End request called", stubConnection.endRequestCalled);
+   }
+}