001/**
002 *
003 * Licensed to the Apache Software Foundation (ASF) under one or more
004 * contributor license agreements.  See the NOTICE file distributed with
005 * this work for additional information regarding copyright ownership.
006 * The ASF licenses this file to You under the Apache License, Version 2.0
007 * (the "License"); you may not use this file except in compliance with
008 * the License.  You may obtain a copy of the License at
009 *
010 *     http://www.apache.org/licenses/LICENSE-2.0
011 *
012 *  Unless required by applicable law or agreed to in writing, software
013 *  distributed under the License is distributed on an "AS IS" BASIS,
014 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
015 *  See the License for the specific language governing permissions and
016 *  limitations under the License.
017 */
018package org.apache.commons.dbcp2.managed;
019
020import java.sql.Connection;
021import java.sql.SQLException;
022import java.util.Map;
023import java.util.Objects;
024import java.util.WeakHashMap;
025
026import javax.transaction.SystemException;
027import javax.transaction.Transaction;
028import javax.transaction.TransactionManager;
029import javax.transaction.TransactionSynchronizationRegistry;
030import javax.transaction.xa.XAResource;
031
032import org.apache.commons.dbcp2.DelegatingConnection;
033
034/**
035 * TransactionRegistry tracks Connections and XAResources in a transacted environment for a single XAConnectionFactory.
036 * <p>
037 * The TransactionRegistry hides the details of transaction processing from the existing DBCP pooling code, and gives
038 * the ManagedConnection a way to enlist connections in a transaction, allowing for the maximal rescue of DBCP.
039 * </p>
040 *
041 * @since 2.0
042 */
043public class TransactionRegistry {
044    private final TransactionManager transactionManager;
045    private final Map<Transaction, TransactionContext> caches = new WeakHashMap<>();
046    private final Map<Connection, XAResource> xaResources = new WeakHashMap<>();
047    private final TransactionSynchronizationRegistry transactionSynchronizationRegistry;
048
049    /**
050     * Creates a TransactionRegistry for the specified transaction manager.
051     *
052     * @param transactionManager
053     *            the transaction manager used to enlist connections.
054     * @param transactionSynchronizationRegistry
055     *              The optional TSR to register synchronizations with
056     * @since 2.6.0
057     */
058    public TransactionRegistry(final TransactionManager transactionManager, final TransactionSynchronizationRegistry transactionSynchronizationRegistry) {
059        this.transactionManager = transactionManager;
060        this.transactionSynchronizationRegistry = transactionSynchronizationRegistry;
061    }
062
063    /**
064     * Provided for backwards compatibility
065     * @param transactionManager the transaction manager used to enlist connections
066     */
067    public TransactionRegistry(final TransactionManager transactionManager) {
068        this (transactionManager, null);
069    }
070
071    /**
072     * Registers the association between a Connection and a XAResource. When a connection is enlisted in a transaction,
073     * it is actually the XAResource that is given to the transaction manager.
074     *
075     * @param connection
076     *            The JDBC connection.
077     * @param xaResource
078     *            The XAResource which managed the connection within a transaction.
079     */
080    public synchronized void registerConnection(final Connection connection, final XAResource xaResource) {
081        Objects.requireNonNull(connection, "connection is null");
082        Objects.requireNonNull(xaResource, "xaResource is null");
083        xaResources.put(connection, xaResource);
084    }
085
086    /**
087     * Gets the XAResource registered for the connection.
088     *
089     * @param connection
090     *            the connection
091     * @return The XAResource registered for the connection; never null.
092     * @throws SQLException
093     *             Thrown when the connection does not have a registered XAResource.
094     */
095    public synchronized XAResource getXAResource(final Connection connection) throws SQLException {
096        Objects.requireNonNull(connection, "connection is null");
097        final Connection key = getConnectionKey(connection);
098        final XAResource xaResource = xaResources.get(key);
099        if (xaResource == null) {
100            throw new SQLException("Connection does not have a registered XAResource " + connection);
101        }
102        return xaResource;
103    }
104
105    /**
106     * Gets the active TransactionContext or null if not Transaction is active.
107     *
108     * @return The active TransactionContext or null if no Transaction is active.
109     * @throws SQLException
110     *             Thrown when an error occurs while fetching the transaction.
111     */
112    public TransactionContext getActiveTransactionContext() throws SQLException {
113        Transaction transaction = null;
114        try {
115            transaction = transactionManager.getTransaction();
116
117            // was there a transaction?
118            if (transaction == null) {
119                return null;
120            }
121
122            // This is the transaction on the thread so no need to check it's status - we should try to use it and
123            // fail later based on the subsequent status
124        } catch (final SystemException e) {
125            throw new SQLException("Unable to determine current transaction ", e);
126        }
127
128        // register the context (or create a new one)
129        synchronized (this) {
130            TransactionContext cache = caches.get(transaction);
131            if (cache == null) {
132                cache = new TransactionContext(this, transaction, transactionSynchronizationRegistry);
133                caches.put(transaction, cache);
134            }
135            return cache;
136        }
137    }
138
139    /**
140     * Unregisters a destroyed connection from {@link TransactionRegistry}.
141     *
142     * @param connection
143     *            A destroyed connection from {@link TransactionRegistry}.
144     */
145    public synchronized void unregisterConnection(final Connection connection) {
146        final Connection key = getConnectionKey(connection);
147        xaResources.remove(key);
148    }
149
150    private Connection getConnectionKey(final Connection connection) {
151        Connection result;
152        if (connection instanceof DelegatingConnection) {
153            result = ((DelegatingConnection<?>) connection).getInnermostDelegateInternal();
154        } else {
155            result = connection;
156        }
157        return result;
158    }
159}