cn.uncode.dal.internal.shards.transaction.MultiDataSourcesTransactionManager.java Source code

Introduction

Here is the source code for cn.uncode.dal.internal.shards.transaction.MultiDataSourcesTransactionManager.java

Source

/*
 * @(#)MultiDataSourcesTransactionManager.java 2012-8-1 ?10:00:00
 *
 * Copyright (c) 2011-2012 Makersoft.org all rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 *
 */
package cn.uncode.dal.internal.shards.transaction;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

import javax.sql.DataSource;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.jdbc.datasource.DataSourceTransactionManager;
import org.springframework.transaction.PlatformTransactionManager;
import org.springframework.transaction.TransactionDefinition;
import org.springframework.transaction.TransactionException;
import org.springframework.transaction.TransactionStatus;
import org.springframework.transaction.support.DefaultTransactionDefinition;
import org.springframework.transaction.support.TransactionSynchronizationManager;
import org.springframework.util.Assert;

public class MultiDataSourcesTransactionManager implements PlatformTransactionManager, InitializingBean {

    private final Log log = LogFactory.getLog(getClass());

    /* ?? */
    private List<DataSource> dataSources;

    private Map<DataSource, DataSourceTransactionManager> transactionManagers = new HashMap<DataSource, DataSourceTransactionManager>();

    /**
     * ??
     */
    private AtomicInteger commitCount = new AtomicInteger(0);

    /**
     * 
     */
    private AtomicInteger rollbackCount = new AtomicInteger(0);

    public void setDataSources(final List<DataSource> dataSources) {
        this.dataSources = dataSources;
    }

    @Override
    public void afterPropertiesSet() throws Exception {
        Assert.notNull(dataSources, "data source can not be null.");

        for (DataSource dataSource : dataSources) {
            DataSourceTransactionManager txManager = new DataSourceTransactionManager(dataSource);
            transactionManagers.put(dataSource, txManager);
        }
        this.addShutdownHook();
    }

    private void addShutdownHook() {
        Runtime.getRuntime().addShutdownHook(new Thread() {
            @SuppressWarnings("static-access")
            @Override
            public void run() {
                // rollback  commit?0
                while (commitCount.get() != 0) {
                    log.info("Waiting for commit transaction.");
                    try {
                        Thread.currentThread().sleep(1);
                    } catch (InterruptedException e) {
                        log.warn("interrupted when shuting down the query executor:\n{}", e);
                    }
                }
                while (rollbackCount.get() != 0) {
                    log.info("Waiting for rollback transaction.");
                    try {
                        Thread.currentThread().sleep(1);
                    } catch (InterruptedException e) {
                        log.warn("interrupted when shuting down the query executor:\n{}", e);
                    }
                }
                log.info("Transaction success.");
            }
        });

    }

    @Override
    public TransactionStatus getTransaction(TransactionDefinition definition) throws TransactionException {

        MultiDataSourcesTransactionStatus transactionStatus = new MultiDataSourcesTransactionStatus();

        log.debug("Operation '" + definition.getName() + "' starting transaction.");

        for (DataSource dataSource : dataSources) {
            DefaultTransactionDefinition defaultTransactionDefinition = new DefaultTransactionDefinition(
                    definition);
            defaultTransactionDefinition.setName(definition.getName());

            PlatformTransactionManager txManager = this.transactionManagers.get(dataSource);
            TransactionStatus status = txManager.getTransaction(defaultTransactionDefinition);

            TransactionSynchronizationManager.setCurrentTransactionName(defaultTransactionDefinition.getName());

            transactionStatus.put(dataSource, status);
        }

        return transactionStatus;

    }

    @Override
    public void commit(TransactionStatus status) throws TransactionException {

        Throwable ex = null;
        //      Collections.reverse(dataSources);
        for (int i = dataSources.size() - 1; i >= 0; i--) {
            DataSource dataSource = dataSources.get(i);
            try {
                commitCount.addAndGet(1);

                log.debug("Committing JDBC transaction");

                DataSourceTransactionManager txManager = this.transactionManagers.get(dataSource);

                TransactionStatus transactionStatus = ((MultiDataSourcesTransactionStatus) status).get(dataSource);
                txManager.commit(transactionStatus);

                log.debug("Commit JDBC transaction success");
            } catch (Throwable e) {
                log.debug("Could not commit JDBC transaction", e);
                ex = e;
            } finally {
                commitCount.addAndGet(-1);
            }
        }

        if (ex != null) {
            throw new RuntimeException(ex);
        }

    }

    @Override
    public void rollback(TransactionStatus status) throws TransactionException {

        Throwable ex = null;

        //Cannot deactivate transaction synchronization - not active
        //      Collections.reverse(dataSources);

        for (int i = dataSources.size() - 1; i >= 0; i--) {
            DataSource dataSource = dataSources.get(i);
            try {
                log.debug("Rolling back JDBC transaction");

                rollbackCount.addAndGet(1);

                DataSourceTransactionManager txManager = this.transactionManagers.get(dataSource);
                TransactionStatus currentStatus = ((MultiDataSourcesTransactionStatus) status).get(dataSource);

                txManager.rollback(currentStatus);

                log.info("Roll back JDBC transaction success");
            } catch (Throwable e) {
                log.info("Could not roll back JDBC transaction", e);
                ex = e;
            } finally {
                rollbackCount.addAndGet(-1);
            }
        }

        if (ex != null) {
            throw new RuntimeException(ex);
        }

    }

}