org.vader.common.spring.TransactionScope.java Source code

Introduction

Here is the source code for org.vader.common.spring.TransactionScope.java

Source

/*
 * Copyright 2014. Vadim Baranov
 *
 *    Licensed under the Apache License, Version 2.0 (the "License");
 *    you may not use this file except in compliance with the License.
 *    You may obtain a copy of the License at
 *
 *        http://www.apache.org/licenses/LICENSE-2.0
 *
 *    Unless required by applicable law or agreed to in writing, software
 *    distributed under the License is distributed on an "AS IS" BASIS,
 *    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *    See the License for the specific language governing permissions and
 *    limitations under the License.
 */
package org.vader.common.spring;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.ObjectFactory;
import org.springframework.beans.factory.config.Scope;
import org.springframework.core.NamedThreadLocal;
import org.springframework.transaction.support.TransactionSynchronizationAdapter;
import org.springframework.transaction.support.TransactionSynchronizationManager;

import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Deque;
import java.util.HashMap;
import java.util.Map;

/**
 * @author Vadim Baranov
 */
public class TransactionScope implements Scope {
    private static final Logger LOG = LoggerFactory.getLogger(TransactionScope.class);

    private static final ThreadLocal<Deque<Map<String, ScopeEntry>>> SCOPES = new NamedThreadLocal<Deque<Map<String, ScopeEntry>>>(
            TransactionScope.class.getName()) {
        @Override
        protected Deque<Map<String, ScopeEntry>> initialValue() {
            final Deque<Map<String, ScopeEntry>> result = new ArrayDeque<>(2);
            result.push(new HashMap<String, ScopeEntry>());
            return result;
        }
    };

    private NonTransactionalBehaviour nonTransactionalBehaviour = NonTransactionalBehaviour.NOT_ALLOWED;

    public NonTransactionalBehaviour getNonTransactionalBehaviour() {
        return nonTransactionalBehaviour;
    }

    public void setNonTransactionalBehaviour(NonTransactionalBehaviour nonTransactionalBehaviour) {
        this.nonTransactionalBehaviour = nonTransactionalBehaviour;
    }

    @Override
    public Object get(String name, ObjectFactory<?> objectFactory) {
        LOG.debug("Get bean={} from transaction scope", name);
        if (!TransactionSynchronizationManager.isSynchronizationActive()) {
            switch (nonTransactionalBehaviour) {
            case PROTOTYPE:
                return objectFactory.getObject();
            case NOT_ALLOWED:
            default:
                throw new IllegalStateException(String.format(
                        "Unable to create bean=%s within transaction scope. No active transaction found", name));
            }
        }
        final ScopeEntry entry = getOrCreateScopeEntry(name);
        if (entry.getBean() != null) {
            LOG.debug("Returns existing bean={} from transaction scope", name);
            return entry.getBean();
        }
        LOG.debug("Creates new bean={} from transaction scope", name);
        final Object newBean = objectFactory.getObject();
        getCurrentScope().get(name).setBean(newBean);
        return newBean;
    }

    @Override
    public Object remove(String name) {
        if (!TransactionSynchronizationManager.isSynchronizationActive()) {
            switch (nonTransactionalBehaviour) {
            case PROTOTYPE:
                break;
            case NOT_ALLOWED:
            default:
                throw new IllegalStateException(String.format(
                        "Unable to remove bean=%s within transaction scope. No active transaction found", name));
            }
        }
        return getCurrentScope().remove(name);
    }

    @Override
    public void registerDestructionCallback(String name, Runnable callback) {
        if (!TransactionSynchronizationManager.isSynchronizationActive()) {
            switch (nonTransactionalBehaviour) {
            case PROTOTYPE:
                break;
            case NOT_ALLOWED:
            default:
                throw new IllegalStateException(
                        String.format("Unable register destruction callback for bean=%s within transaction scope. "
                                + "No active transaction found", name));
            }
        }
        getOrCreateScopeEntry(name).getCallbacks().add(callback);
    }

    @Override
    public Object resolveContextualObject(String key) {
        return null;
    }

    @Override
    public String getConversationId() {
        return TransactionSynchronizationManager.getCurrentTransactionName();
    }

    private static Map<String, ScopeEntry> getCurrentScope() {
        return SCOPES.get().peek();
    }

    private static ScopeEntry getOrCreateScopeEntry(String name) {
        final Map<String, ScopeEntry> scope = getCurrentScope();
        if (scope.isEmpty()) {
            registerSynchronization();
        }
        final ScopeEntry entry = scope.get(name);
        if (entry != null) {
            return entry;
        }
        final ScopeEntry newEntry = new ScopeEntry();
        scope.put(name, newEntry);
        return newEntry;
    }

    private static void registerSynchronization() {
        LOG.debug("Registers transaction synchronization");
        TransactionSynchronizationManager.registerSynchronization(new TransactionSynchronizationAdapter() {
            @Override
            public void suspend() {
                LOG.debug("Push new transaction scope");
                SCOPES.get().push(new HashMap<String, ScopeEntry>());
            }

            @Override
            public void resume() {
                LOG.debug("Pop old transaction scope");
                SCOPES.get().pop();
            }

            @Override
            public void afterCompletion(int status) {
                final Map<String, ScopeEntry> scope = getCurrentScope();
                LOG.debug("Destroying {} beans", scope.size());
                for (Map.Entry<String, ScopeEntry> entry : scope.entrySet()) {
                    for (Runnable runnable : entry.getValue().getCallbacks()) {
                        LOG.debug("Executes destruction callback for bean={}", entry.getKey());
                        runnable.run();
                    }
                }
                scope.clear();
                LOG.debug("Destruction completed");
            }
        });
    }

    /**
     * Scope entry.
     */
    private static class ScopeEntry {
        private Object bean;
        private Collection<Runnable> callbacks = new ArrayList<>();

        public Object getBean() {
            return bean;
        }

        public void setBean(Object bean) {
            this.bean = bean;
        }

        public Collection<Runnable> getCallbacks() {
            return callbacks;
        }

        public void setCallbacks(Collection<Runnable> callbacks) {
            this.callbacks = callbacks;
        }
    }
}