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     * Provided for backwards compatibility
051     * @param transactionManager the transaction manager used to enlist connections
052     */
053    public TransactionRegistry(final TransactionManager transactionManager) {
054        this (transactionManager, null);
055    }
056
057    /**
058     * Creates a TransactionRegistry for the specified transaction manager.
059     *
060     * @param transactionManager
061     *            the transaction manager used to enlist connections.
062     * @param transactionSynchronizationRegistry
063     *              The optional TSR to register synchronizations with
064     * @since 2.6.0
065     */
066    public TransactionRegistry(final TransactionManager transactionManager, final TransactionSynchronizationRegistry transactionSynchronizationRegistry) {
067        this.transactionManager = transactionManager;
068        this.transactionSynchronizationRegistry = transactionSynchronizationRegistry;
069    }
070
071    /**
072     * Gets the active TransactionContext or null if not Transaction is active.
073     *
074     * @return The active TransactionContext or null if no Transaction is active.
075     * @throws SQLException
076     *             Thrown when an error occurs while fetching the transaction.
077     */
078    public TransactionContext getActiveTransactionContext() throws SQLException {
079        Transaction transaction = null;
080        try {
081            transaction = transactionManager.getTransaction();
082
083            // was there a transaction?
084            if (transaction == null) {
085                return null;
086            }
087
088            // This is the transaction on the thread so no need to check it's status - we should try to use it and
089            // fail later based on the subsequent status
090        } catch (final SystemException e) {
091            throw new SQLException("Unable to determine current transaction ", e);
092        }
093
094        // register the context (or create a new one)
095        synchronized (this) {
096            TransactionContext cache = caches.get(transaction);
097            if (cache == null) {
098                cache = new TransactionContext(this, transaction, transactionSynchronizationRegistry);
099                caches.put(transaction, cache);
100            }
101            return cache;
102        }
103    }
104
105    private Connection getConnectionKey(final Connection connection) {
106        final Connection result;
107        if (connection instanceof DelegatingConnection) {
108            result = ((DelegatingConnection<?>) connection).getInnermostDelegateInternal();
109        } else {
110            result = connection;
111        }
112        return result;
113    }
114
115    /**
116     * Gets the XAResource registered for the connection.
117     *
118     * @param connection
119     *            the connection
120     * @return The XAResource registered for the connection; never null.
121     * @throws SQLException
122     *             Thrown when the connection does not have a registered XAResource.
123     */
124    public synchronized XAResource getXAResource(final Connection connection) throws SQLException {
125        Objects.requireNonNull(connection, "connection is null");
126        final Connection key = getConnectionKey(connection);
127        final XAResource xaResource = xaResources.get(key);
128        if (xaResource == null) {
129            throw new SQLException("Connection does not have a registered XAResource " + connection);
130        }
131        return xaResource;
132    }
133
134    /**
135     * Registers the association between a Connection and a XAResource. When a connection is enlisted in a transaction,
136     * it is actually the XAResource that is given to the transaction manager.
137     *
138     * @param connection
139     *            The JDBC connection.
140     * @param xaResource
141     *            The XAResource which managed the connection within a transaction.
142     */
143    public synchronized void registerConnection(final Connection connection, final XAResource xaResource) {
144        Objects.requireNonNull(connection, "connection is null");
145        Objects.requireNonNull(xaResource, "xaResource is null");
146        xaResources.put(connection, xaResource);
147    }
148
149    /**
150     * Unregisters a destroyed connection from {@link TransactionRegistry}.
151     *
152     * @param connection
153     *            A destroyed connection from {@link TransactionRegistry}.
154     */
155    public synchronized void unregisterConnection(final Connection connection) {
156        final Connection key = getConnectionKey(connection);
157        xaResources.remove(key);
158    }
159}