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 *      http://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}