Merge pull request #50 from sshankar/customize_con_on_lease_release

Add ability to customize connection on lease and return to pool.

btm/src/main/java/bitronix/tm/resource/jdbc/ConnectionCustomizer.java

@@ -31,6 +31,20 @@ public interface ConnectionCustomizer {
      * @param uniqueName the PoolingDataSource unique name.
      */
     public void onAcquire(Connection connection, String uniqueName);
+
+    /**
+     * Called when the physical connection is leased from the pool.
+     * @param connection the physical connection.
+     * @param uniqueName the PoolingDataSource unique name.
+     */
+    public void onLease(Connection connection, String uniqueName);
+
+    /**
+     * Called when the physical connection is returned to the pool.
+     * @param connection the physical connection.
+     * @param uniqueName the PoolingDataSource unique name.
+     */
+    public void onRelease(Connection connection, String uniqueName);
 
     /**
      * Called when the physical connection is destroyed.

btm/src/main/java/bitronix/tm/resource/jdbc/JdbcPooledConnection.java

@@ -226,6 +226,7 @@ public boolean release() throws SQLException {
             // Only requeue a connection if it is no longer in use.  In the case of non-shared connections,
             // usageCount will always be 0 here, so the default behavior is unchanged.
             if (usageCount == 0) {
+                poolingDataSource.fireOnRelease(connection);
                 try {
                     TransactionContextHelper.requeue(this, poolingDataSource);
                 } catch (BitronixSystemException ex) {
@@ -309,6 +310,8 @@ public Object getConnectionHandle() throws Exception {
         }
 
         if (log.isDebugEnabled()) { log.debug("got connection handle from " + this); }
+        poolingDataSource.fireOnLease(connection);
+
         return getConnectionHandle(connection);
     }
 

btm/src/main/java/bitronix/tm/resource/jdbc/PoolingDataSource.java

@@ -279,6 +279,26 @@ void fireOnAcquire(Connection connection) {
         }
     }
 
+    void fireOnLease(Connection connection){
+        for (ConnectionCustomizer connectionCustomizer : connectionCustomizers) {
+            try {
+                connectionCustomizer.onLease(connection, getUniqueName());
+            } catch (Exception ex){
+                log.warn("ConnectionCustomizer.onLease() failed for " + connectionCustomizer, ex);
+            }
+        }
+    }
+
+    void fireOnRelease(Connection connection){
+        for (ConnectionCustomizer connectionCustomizer : connectionCustomizers) {
+            try {
+                connectionCustomizer.onRelease(connection, getUniqueName());
+            } catch (Exception ex){
+                log.warn("ConnectionCustomizer.onRelease() failed for " + connectionCustomizer, ex);
+            }
+        }
+    }
+
     void fireOnDestroy(Connection connection) {
         for (ConnectionCustomizer connectionCustomizer : connectionCustomizers) {
             try {
@@ -289,7 +309,6 @@ void fireOnDestroy(Connection connection) {
         }
     }
 
-
     /* Implementation of DataSource interface */
     @Override
     public Connection getConnection() throws SQLException {

btm/src/test/java/bitronix/tm/resource/jdbc/ConnectionCustomizerTest.java

@@ -0,0 +1,62 @@
+/*
+ * Copyright (C) 2006-2013 Bitronix Software (http://www.bitronix.be)
+ * 
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with 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 bitronix.tm.resource.jdbc;
+
+import bitronix.tm.resource.jdbc.PoolingDataSource;
+import bitronix.tm.mock.resource.jdbc.MockitoXADataSource;
+
+import java.sql.Connection;
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.InOrder;
+
+import static org.mockito.Mockito.*;
+
+public class ConnectionCustomizerTest {
+    private ConnectionCustomizer customizer;
+
+    @Before public void setup(){
+        customizer = mock(ConnectionCustomizer.class);
+    }
+
+    @Test public void testCustomizerCall() throws Exception{
+        String name = "pds";
+        PoolingDataSource pds = new PoolingDataSource();
+        pds.setUniqueName(name);
+        pds.setXaDataSource(new MockitoXADataSource());
+        pds.setMinPoolSize(1);
+        pds.setMaxPoolSize(1);
+        pds.setXaDataSource(new MockitoXADataSource());
+        pds.addConnectionCustomizer(customizer);
+        pds.init();
+
+        Connection connection = pds.getConnection(); // onAcquire, onLease
+        connection.close(); // onRelease
+
+        connection = pds.getConnection(); // onLease
+        connection.close(); // onRelease
+
+        pds.close(); // onDestroy
+
+        InOrder callOrder = inOrder(customizer);
+        callOrder.verify(customizer).onAcquire(any(Connection.class), eq(name));
+        callOrder.verify(customizer).onLease(any(Connection.class), eq(name));
+        callOrder.verify(customizer).onRelease(any(Connection.class), eq(name));
+        callOrder.verify(customizer).onLease(any(Connection.class), eq(name));
+        callOrder.verify(customizer).onRelease(any(Connection.class), eq(name));
+        callOrder.verify(customizer).onDestroy(any(Connection.class), eq(name));
+    }
+}