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