ServiceExporterBeanPostProcessor

package com.atlassian.plugin.spring.scanner.runtime.impl;

import com.atlassian.plugin.spring.scanner.annotation.export.ExportAsDevService;
import com.atlassian.plugin.spring.scanner.annotation.export.ExportAsService;
import com.atlassian.plugin.spring.scanner.annotation.export.ModuleType;
import com.atlassian.plugin.spring.scanner.util.CommonConstants;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Iterables;
import org.osgi.framework.Bundle;
import org.osgi.framework.BundleContext;
import org.osgi.framework.ServiceRegistration;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.aop.support.AopUtils;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
import org.springframework.beans.factory.config.DestructionAwareBeanPostProcessor;

import java.lang.annotation.Annotation;
import java.util.Hashtable;
import java.util.List;

import static com.atlassian.plugin.spring.scanner.runtime.impl.util.AnnotationIndexReader.getIndexFilesForProfiles;
import static com.atlassian.plugin.spring.scanner.runtime.impl.util.AnnotationIndexReader.readAllIndexFilesForProduct;
import static com.atlassian.plugin.spring.scanner.runtime.impl.util.AnnotationIndexReader.splitProfiles;

/**
 * A BeanPostProcessor that exports OSGi services for beans annotated with both a *Component annotation and the
 * ExportAsService annotation This essentially does the same thing as the "public=true" on an atlassian-plugin.xml
 * component entry.
 * <p>
 * This is implemented as a BeanPostProcessor because we need to service to come and go as the bean is
 * created/destroyed
 */
public class ServiceExporterBeanPostProcessor implements DestructionAwareBeanPostProcessor, InitializingBean {
    public static final String OSGI_SERVICE_SUFFIX = "_osgiService";
    static final String ATLASSIAN_DEV_MODE_PROP = "atlassian.dev.mode";

    private static final Logger log = LoggerFactory.getLogger(ServiceExporterBeanPostProcessor.class);

    private final boolean isDevMode = Boolean.parseBoolean(System.getProperty(ATLASSIAN_DEV_MODE_PROP, "false"));
    private final BundleContext bundleContext;
    private ConfigurableListableBeanFactory beanFactory;
    private String profileName;
    private final ExportedSeviceManager serviceManager;

    private ImmutableMap<String, String[]> exports;

    public ServiceExporterBeanPostProcessor(final BundleContext bundleContext, final ConfigurableListableBeanFactory beanFactory) {
        this(bundleContext, beanFactory, new ExportedSeviceManager());
    }

    ServiceExporterBeanPostProcessor(final BundleContext bundleContext, final ConfigurableListableBeanFactory beanFactory,
                                     ExportedSeviceManager serviceManager) {
        this.bundleContext = bundleContext;
        this.beanFactory = beanFactory;
        this.profileName = null;
        this.serviceManager = serviceManager;
    }

    // Used to inject profileName
    @SuppressWarnings("UnusedDeclaration")
    public void setProfileName(final String profileName) {
        this.profileName = profileName;
    }

    @Override
    public void afterPropertiesSet() throws Exception {
        final ImmutableMap.Builder<String, String[]> exportBuilder = ImmutableMap.builder();

        final Bundle bundle = bundleContext.getBundle();

        final String[] profileNames = splitProfiles(profileName);
        parseExportsForExportFile(exportBuilder, CommonConstants.COMPONENT_EXPORT_KEY, profileNames, bundle);
        if (isDevMode) {
            parseExportsForExportFile(exportBuilder, CommonConstants.COMPONENT_DEV_EXPORT_KEY, profileNames, bundle);
        }
        exports = exportBuilder.build();
    }

    private void parseExportsForExportFile(
            final ImmutableMap.Builder<String, String[]> exportBuilder, final String exportFileName,
            final String[] profileNames, final Bundle bundle) {
        final String[] defaultInterfaces = {};
        for (final String fileToRead : getIndexFilesForProfiles(profileNames, exportFileName)) {
            final List<String> exportData = readAllIndexFilesForProduct(fileToRead, bundle, bundleContext);
            for (final String export : exportData) {
                final String[] targetAndInterfaces = export.split("#");
                final String target = targetAndInterfaces[0];
                final String[] interfaces = (targetAndInterfaces.length > 1)
                        ? targetAndInterfaces[1].split(",")
                        : defaultInterfaces;
                exportBuilder.put(target, interfaces);
            }
        }
    }

    @Override
    public void postProcessBeforeDestruction(final Object bean, final String beanName) throws BeansException {
        if (serviceManager.hasService(bean)) {
            serviceManager.unregisterService(bundleContext, bean);

            final String serviceName = getServiceName(beanName);
            if (beanFactory.containsBean(serviceName)) {
                final Object serviceBean = beanFactory.getBean(serviceName);

                if (null != serviceBean) {
                    beanFactory.destroyBean(serviceName, serviceBean);
                }
            }
        }
    }

    @Override
    public Object postProcessBeforeInitialization(final Object bean, final String beanName) throws BeansException {
        return bean;
    }

    @Override
    public Object postProcessAfterInitialization(final Object bean, final String beanName) throws BeansException {
        Class<?>[] interfaces = {};
        // If the actual bean is some proxy object because of Spring AOP, this extracts the actual declared class of the bean.
        // The answer will only be different from getClass() when AOP is in play - in such case this will retrieve
        // the real class used in the code ignoring AOP proxy
        final Class<?> beanTargetClass = AopUtils.getTargetClass(bean);
        final String beanClassName = beanTargetClass.getName();

        if (exports.containsKey(beanClassName) || isPublicComponent(beanTargetClass)) {
            if (exports.containsKey(beanClassName)) {
                // We need to turn the string interface names from the exports file into actual Class objects. We do this by
                // asking the bean's class loader for each interface by name. We could walking the class hierarchy looking for
                // the ones we want, but that seems a bit more roundabout.
                final ImmutableList.Builder<Class<?>> interfaceBuilder = ImmutableList.builder();
                final ClassLoader beanClassLoader = bean.getClass().getClassLoader();
                final String[] interfaceNames = exports.get(beanClassName);
                for (final String interfaceName : interfaceNames) {
                    try {
                        final Class interfaceClass = beanClassLoader.loadClass(interfaceName);
                        interfaceBuilder.add(interfaceClass);
                    } catch (final ClassNotFoundException ecnf) {
                        log.warn("Cannot find class for export '" + interfaceName + "' of bean '" + beanName + "': " + ecnf);
                        // And drop it - there's not much else we can do
                    }
                }
                interfaces = Iterables.toArray(interfaceBuilder.build(), Class.class);
            } else if (hasAnnotation(beanTargetClass, ModuleType.class)) {
                interfaces = beanTargetClass.getAnnotation(ModuleType.class).value();
            } else if (hasAnnotation(beanTargetClass, ExportAsService.class)) {
                interfaces = beanTargetClass.getAnnotation(ExportAsService.class).value();
            } else if (hasAnnotation(beanTargetClass, ExportAsDevService.class)) {
                interfaces = beanTargetClass.getAnnotation(ExportAsDevService.class).value();
            }

            //if they didn't specify any interfaces, calculate them
            if (interfaces.length < 1) {
                // This only gets direct interfaces, not inherited ones, which is documented behaviour.
                // It's desireable because it reduces brittleness with respect to superclass changes.
                interfaces = beanTargetClass.getInterfaces();

                //if we still don't have any, just export with the classname (yes, OSGi allows this.
                if (interfaces.length < 1) {
                    interfaces = new Class<?>[]{beanTargetClass};
                }
            }

            try {
                final ServiceRegistration serviceRegistration = serviceManager.registerService(
                        bundleContext, bean, beanName, new Hashtable<String, Object>(), interfaces);
                final String serviceName = getServiceName(beanName);
                beanFactory.initializeBean(serviceRegistration, serviceName);
            } catch (final Exception e) {
                log.error("Unable to register bean '" + beanName + "' as an OSGi exported service", e);
            }
        }

        return bean;
    }

    private boolean isPublicComponent(final Class beanTargetClass) {
        return hasAnnotation(beanTargetClass, ModuleType.class)
                || hasAnnotation(beanTargetClass, ExportAsService.class)
                || (hasAnnotation(beanTargetClass, ExportAsDevService.class) && isDevMode);
    }

    private boolean hasAnnotation(final Class beanTargetClass, final Class<? extends Annotation> annotationClass) {
        return beanTargetClass.isAnnotationPresent(annotationClass);
    }

    private String getServiceName(final String beanName) {
        return beanName + OSGI_SERVICE_SUFFIX;
    }

}