/atlassian-plugins-osgi/src/main/java/com/atlassian/plugin/osgi/hostcomponents/impl/DefaultComponentRegistrar.java

package com.atlassian.plugin.osgi.hostcomponents.impl;

import com.atlassian.plugin.hostcontainer.HostContainer;
import com.atlassian.plugin.osgi.hostcomponents.ComponentRegistrar;
import com.atlassian.plugin.osgi.hostcomponents.ContextClassLoaderStrategy;
import com.atlassian.plugin.osgi.hostcomponents.HostComponentRegistration;
import com.atlassian.plugin.osgi.hostcomponents.InstanceBuilder;
import com.atlassian.plugin.osgi.hostcomponents.PropertyBuilder;
import com.atlassian.plugin.util.ContextClassLoaderSettingInvocationHandler;
import com.google.common.base.Function;
import com.google.common.collect.Maps;
import org.osgi.framework.Bundle;
import org.osgi.framework.BundleContext;
import org.osgi.framework.ServiceFactory;
import org.osgi.framework.ServiceRegistration;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.lang.reflect.InvocationHandler;
import java.lang.reflect.Method;
import java.lang.reflect.Proxy;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CopyOnWriteArrayList;

/**
 * Default component registrar that also can write registered host components into the OSGi service registry.
 */
public class DefaultComponentRegistrar implements ComponentRegistrar {
    private static final Logger log = LoggerFactory.getLogger(DefaultComponentRegistrar.class);

    private final List<HostComponentRegistration> registry = new CopyOnWriteArrayList<HostComponentRegistration>();

    public InstanceBuilder register(final Class<?>... mainInterfaces) {
        final Registration reg = new Registration(mainInterfaces);
        registry.add(reg);
        return new DefaultInstanceBuilder(reg);
    }

    public List<ServiceRegistration> writeRegistry(final BundleContext ctx) {
        final ArrayList<ServiceRegistration> services = new ArrayList<ServiceRegistration>();

        for (final HostComponentRegistration reg : new ArrayList<HostComponentRegistration>(registry)) {
            if (Arrays.asList(reg.getMainInterfaceClasses()).contains(HostContainer.class)) {
                log.warn("Cannot register a HostContainer as a host component, skipping");
                registry.remove(reg);
                continue;
            }

            final String[] names = reg.getMainInterfaces();

            reg.getProperties().put(HOST_COMPONENT_FLAG, Boolean.TRUE.toString());

            // If no bean name specified, generate one that will be consistent across restarts
            final String beanName = reg.getProperties().get(PropertyBuilder.BEAN_NAME);
            if (beanName == null) {
                String genKey = String.valueOf(Arrays.asList(reg.getMainInterfaces()).hashCode());
                reg.getProperties().put(PropertyBuilder.BEAN_NAME, "hostComponent-" + genKey);
            }

            if (log.isDebugEnabled()) {
                log.debug("Registering: " + Arrays.asList(names) + " instance " + reg.getInstance() + "with properties: " + reg.getProperties());
            }

            if (names.length == 0) {
                log.warn("Host component " + beanName + " of instance " + reg.getInstance() + " has no interfaces");
            }

            Object service = reg.getInstance();

            if (!ContextClassLoaderStrategy.USE_PLUGIN.name().equals(reg.getProperties().get(PropertyBuilder.CONTEXT_CLASS_LOADER_STRATEGY))) {
                service = createContextClassLoaderSettingProxy(reg.getMainInterfaceClasses(), service);
            }

            if (Boolean.parseBoolean(reg.getProperties().get(PropertyBuilder.TRACK_BUNDLE))) {
                service = createTrackBundleProxy(reg.getMainInterfaceClasses(), service);
            }

            final ServiceRegistration sreg = ctx.registerService(names, service, reg.getProperties());
            if (sreg != null) {
                services.add(sreg);
            }
        }
        return Collections.unmodifiableList(services);
    }

    public List<HostComponentRegistration> getRegistry() {
        return Collections.unmodifiableList(registry);
    }

    /**
     * Wraps the service in a dynamic proxy that ensures all methods are executed with the object class's class loader
     * as the context class loader
     *
     * @param interfaces The interfaces to proxy
     * @param service    The instance to proxy
     * @return A proxy that wraps the service
     */
    private Object createContextClassLoaderSettingProxy(final Class<?>[] interfaces, final Object service) {
        final Function<Object, Object> transformer = new Function<Object, Object>() {
            public Object apply(final Object service) {
                return Proxy.newProxyInstance(DefaultComponentRegistrar.class.getClassLoader(), interfaces,
                        new ContextClassLoaderSettingInvocationHandler(service));
            }
        };

        if (!(service instanceof ServiceFactory)) {
            return transformer.apply(service);
        }
        return new TransformingServiceFactory((ServiceFactory) service) {
            @Override
            protected Object transform(Bundle bundle, ServiceRegistration registration, Object service) {
                return transformer.apply(service);
            }
        };

    }

    private ServiceFactory createTrackBundleProxy(final Class<?>[] interfaces, final Object service) {
        final ServiceFactory delegate = service instanceof ServiceFactory ?
                (ServiceFactory) service :
                new InstanceServiceFactory(service);

        return new TransformingServiceFactory(delegate) {
            @Override
            protected Object transform(final Bundle bundle, final ServiceRegistration registration, final Object service) {
                return Proxy.newProxyInstance(DefaultComponentRegistrar.class.getClassLoader(), interfaces,
                        new BundleTrackingInvocationHandler(bundle, service));
            }
        };
    }

    private static class BundleTrackingInvocationHandler implements InvocationHandler {

        private final Bundle bundle;
        private final Object service;

        private BundleTrackingInvocationHandler(Bundle bundle, Object service) {
            this.bundle = bundle;
            this.service = service;
        }

        @Override
        public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
            Bundle original = CallingBundleStore.get();
            try {
                CallingBundleStore.set(bundle);
                return method.invoke(service, args);
            } finally {
                CallingBundleStore.set(original);
            }
        }
    }

    private static class InstanceServiceFactory implements ServiceFactory {

        private final Object service;

        private InstanceServiceFactory(final Object service) {
            this.service = service;
        }

        @Override
        public Object getService(final Bundle bundle, final ServiceRegistration registration) {
            return service;
        }

        @Override
        public void ungetService(final Bundle bundle, final ServiceRegistration registration, final Object service) {
            // no-op
        }
    }

    private static abstract class TransformingServiceFactory implements ServiceFactory {

        private final ServiceFactory delegate;
        private final Map<Long, Object> bundleIdToOriginalService;

        private TransformingServiceFactory(ServiceFactory delegate) {
            this.delegate = delegate;
            this.bundleIdToOriginalService = Maps.newConcurrentMap();
        }


        @Override
        public final Object getService(final Bundle bundle, final ServiceRegistration registration) {
            final Object service = delegate.getService(bundle, registration);
            final Object transformed = transform(bundle, registration, service);
            // Store the original service after it has been transformed
            // in case the transformer throws an exception
            bundleIdToOriginalService.put(bundle.getBundleId(), service);
            return transformed;
        }

        @Override
        public final void ungetService(final Bundle bundle, final ServiceRegistration registration, final Object transformed) {
            final Object service = bundleIdToOriginalService.remove(bundle.getBundleId());
            if (service != null) {
                delegate.ungetService(bundle, registration, service);
            }
        }

        protected abstract Object transform(Bundle bundle, ServiceRegistration registration, Object service);

    }

}