DefaultLifecycleManager

package com.atlassian.sal.core.lifecycle;

import com.atlassian.plugin.PluginAccessor;
import com.atlassian.plugin.event.PluginEventListener;
import com.atlassian.plugin.event.PluginEventManager;
import com.atlassian.plugin.event.events.PluginDisablingEvent;
import com.atlassian.plugin.event.events.PluginEnabledEvent;
import com.atlassian.plugin.event.events.PluginFrameworkShuttingDownEvent;
import com.atlassian.plugin.event.events.PluginFrameworkStartedEvent;
import com.atlassian.sal.api.lifecycle.LifecycleAware;
import com.atlassian.sal.api.lifecycle.LifecycleManager;
import org.osgi.framework.Bundle;
import org.osgi.framework.BundleContext;
import org.osgi.framework.ServiceEvent;
import org.osgi.framework.ServiceListener;
import org.osgi.framework.ServiceReference;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.beans.factory.InitializingBean;

import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.function.Consumer;
import java.util.function.Predicate;

import static com.atlassian.plugin.osgi.util.OsgiHeaderUtil.getPluginKey;
import static org.osgi.framework.Constants.OBJECTCLASS;

/**
 * Manage the tracking of {@link LifecycleAware} services registered to OSGi and call their {@link LifecycleAware#onStart}
 * and {@link LifecycleAware#onStop()} when required.
 * <p>
 * The {@code onStart} method is called when the application is setup according to {@link LifecycleManager#isApplicationSetUp}, the
 * plugin is enabled, and the plugin framework has started.
 * <p>
 * The {@code onStop} method is called if {@code onStart} was called, when plugin and it's dependencies are still present.
 * <p>
 * This class has concurrency controls for the case of concurrent service registration/deregistration with plugin and application
 * state activity. It does not account for concurrency of plugin notifications against themselves. This does not happen during
 * normal usage.
 */
public abstract class DefaultLifecycleManager implements LifecycleManager, InitializingBean, DisposableBean {
    private static final Logger log = LoggerFactory.getLogger(DefaultLifecycleManager.class);

    private final PluginEventManager pluginEventManager;
    private final PluginAccessor pluginAccessor;
    private final BundleContext bundleContext;

    /**
     * The LifecycleAware services which we know about via OSGi notifications, but which we have not yet called onStart() on.
     * <p>
     * The documentation for {@link org.osgi.framework.ServiceReference} is quite explicit that while a framework may return
     * multiple ServiceReference objects referring to the same ServiceRegistration, all such ServiceReference objects are equals()
     * and have the same hashCode(), so it is safe to use them in a set like this.
     * <p>
     * Inspection of Felix code (as at 4.2.1 of this writing) shows that in fact Felix is careful to have only a single
     * ServiceReference for each ServiceRegistration. If we move to a framework where this is not so, or if this causes problems,
     * we'd need to replace this with a map from the ServiceReference.getProperty(SERVICE_ID) to ServiceReference so we can key
     * service tracking off SERVICE_ID.
     */
    private final Set<ServiceReference<LifecycleAware>> pendingOnStart;
    private final Set<ServiceReference<LifecycleAware>> pendingOnStop;
    private final ServiceListener serviceListener;

    private boolean started;

    public DefaultLifecycleManager(
            final PluginEventManager pluginEventManager,
            final PluginAccessor pluginAccessor,
            final BundleContext bundleContext) {
        this.pluginEventManager = pluginEventManager;
        this.pluginAccessor = pluginAccessor;
        this.bundleContext = bundleContext;
        this.pendingOnStart = Collections.synchronizedSet(new HashSet<>());
        this.pendingOnStop = Collections.synchronizedSet(new HashSet<>());
        this.serviceListener = new LifecycleAwareServiceListener();
        this.started = false;
    }

    @Override
    public void afterPropertiesSet() throws Exception {
        pluginEventManager.register(this);
        // Gather all the current LifecycleAware instances, and watch for new ones
        final String filter = "(" + OBJECTCLASS + "=" + LifecycleAware.class.getName() + ")";
        // Since OSGi doesn't (afaics) provide an atomic way to add a ServiceListener and get brought up to date, we need to be a
        // bit careful here. If we miss an add, we won't call an onStart() we should, and if we miss a remove we'll possibly leak.
        // So we starting listening, add everything we can see, and discard anything invalid.

        // An an aside, add everything, add listener, add everything doesn't work - a service that is removed after the snapshot
        // for the second add everything is readded when this snapshot is added, and so you end up missing a remove.

        // Services that come and go here before this point are irrelevant
        bundleContext.addServiceListener(serviceListener, filter);
        // Anything added from now on is captured, so we don't lose services
        final Collection<ServiceReference<LifecycleAware>> services =
                bundleContext.getServiceReferences(LifecycleAware.class, null);
        // Our snapshot may contain services that have already made it to pending. This means we might get double adds (which the
        // set takes care of), or we might re-add something in our snapshot that gets removed in this window. Such a service might
        // have turned up before we started listening or after, so the listener might think the remove is spurious or normal, but
        // the service is in our snapshot.
        pendingOnStart.addAll(services);
        // Any remove from now on is correctly processed - we've gathered everything and put it it pendingOnStart. However,
        // we might have re-added some service that was removed post snapshot but before we updated pendingOnStart. To deal
        // with these stale services, we do a sweep and chuck them out.  We can sweep the services from the getServiceReferences
        // snapshot, because those that arrived via the listener will be present when removed. A service added before the listener
        // is added but removed afterwards, but before the snapshot, is just ignored as a spurious remove.
        for (final ServiceReference<LifecycleAware> service : services) {
            // The documentation for getBundle() says this is a sanctioned way to detect deregistration
            if (null == service.getBundle()) {
                pendingOnStart.remove(service);
            }
        }
    }

    @SuppressWarnings("unchecked")
    private void clearPendingOnStop() {
        if (pendingOnStop.isEmpty()) {
            // all onStop() delivered
            return;
        }

        // it's safe to cast here, we know what pendingOnStop is
        log.warn("Failed to notify with LifecycleAware.onStop(): {}",
                LifecycleLog.listPluginKeys(Arrays.asList((ServiceReference<Object>[]) pendingOnStop.toArray())));

        pendingOnStop.clear();
    }

    @Override
    public void destroy() throws Exception {
        bundleContext.removeServiceListener(serviceListener);
        // We  clear out our pendingOnStart, to assist garbage collection. This is just a heuristic - if this code moves
        // this is *not* sufficient to prevent leaks, since a thread could be in serviceListener now, having arrived before we
        // removed the listener but not yet finished. Since this is a @PreDestroy and things are going away, the GC will save us.
        // If we do move this code, we should probably move pendingOnStart to the serviceListener and have their lifetimes
        // match, but I'm not going to pre-emptively solve this synchronization problem without knowledge of the lifetimes and where
        // things are moving to.
        pendingOnStart.clear();
        clearPendingOnStop();
        pluginEventManager.unregister(this);
    }

    @PluginEventListener
    public void onPluginFrameworkStarted(final PluginFrameworkStartedEvent event) {
        startIfApplicationSetup();
    }

    @PluginEventListener
    public void onPluginEnabled(final PluginEnabledEvent event) {
        // I don't think it's actually worth restricting to the "right" plugin here - if a whole lot of stuff is enabled when
        // we get the event, we may as well start them all in one hit. This will stop lots of doubling handling of outstanding
        // LifecycleAware instances. However, this is called a lot before startup, so we hoist the started check.
        if (started) {
            notifyStartableLifecycleAwares();
        }
    }

    @PluginEventListener
    public void onPluginFrameworkShuttingDown(final PluginFrameworkShuttingDownEvent event) {
        // We notifying all waiting onStop() here, as the next thing after this event is
        // all plugins unregistered

        final Collection<ServiceReference<LifecycleAware>> completed =
                notifyLifecycleAwares(pendingOnStop, this::notifyOnStopIfEnabled);

        // Now remove those we handled.
        pendingOnStop.removeAll(completed);
    }

    private Predicate<ServiceReference<LifecycleAware>> notifyIfMyEvent(final PluginDisablingEvent event) {
        return lifecycleAwareServiceReference -> {
            final Bundle bundle = lifecycleAwareServiceReference.getBundle();
            if (bundle == null) {
                log.warn("Discarding onStop() for stale LifecycleAware");
                // This reference is stale, so we report true as it can be discarded
                return true;
            }

            // Notify if the event is for a plugin that waits for onStop()
            return getPluginKey(bundle).equals(event.getPlugin().getKey())
                    && notifyOnStopIfEnabled(lifecycleAwareServiceReference);
        };
    }

    @PluginEventListener
    public void onPluginDisabling(final PluginDisablingEvent event) {
        final Collection<ServiceReference<LifecycleAware>> completed =
                notifyLifecycleAwares(pendingOnStop, notifyIfMyEvent(event));

        // Now remove those we handled.
        pendingOnStop.removeAll(completed);
    }

    public void start() {
        startIfApplicationSetup();
    }

    private void startIfApplicationSetup() {
        final boolean doSetup = !started && isApplicationSetUp();
        if (doSetup) {
            started = true;
            // We could call the internal notify irrespective of doSetup, since it guards internally, but since we need to guard
            // for subclasses, we may as well hoist this in here and save a bit of effort in the no-op case.
            notifyStartableLifecycleAwares();
            notifyOnStart();
        }
    }

    /**
     * An override point for subclasses that require callback notification after the initial start of the system.
     */
    protected void notifyOnStart() {
    }

    private void notifyStartableLifecycleAwares() {
        final Collection<ServiceReference<LifecycleAware>> completed = notifyLifecycleAwares(pendingOnStart,
                this::notifyOnStartIfStartedAndEnabled);

        // Now remove those we handled.
        pendingOnStart.removeAll(completed);
        // and mark them waiting onStop
        pendingOnStop.addAll(completed);
    }

    private Collection<ServiceReference<LifecycleAware>> notifyLifecycleAwares(final Set<ServiceReference<LifecycleAware>> lifeCycleAwares,
                                                                               final Predicate<ServiceReference<LifecycleAware>> event) {
        // Grab a copy to operate on, so we don't call onStart()/onStop() holding a lock. We're using toArray here to keep the
        // synchronization clear at the cost of the unchecked cast below. If we went via ImmutableList.copyOf, we'd need to either
        // synchronize externally or depend on the implementation of copyOf - see Collections.synchronizedSet re iteration.
        final Object[] pending = lifeCycleAwares.toArray();

        // Notify those we can, remembering what we handled.
        final List<ServiceReference<LifecycleAware>> completed = new ArrayList<ServiceReference<LifecycleAware>>(pending.length);
        for (final Object serviceRaw : pending) {
            @SuppressWarnings("unchecked")
            // serviceRaw comes from Object[] pending which is a snapshot of lifeCycleAwares so the cast is safe
            final ServiceReference<LifecycleAware> service = (ServiceReference<LifecycleAware>) serviceRaw;
            if (event.test(service)) {
                completed.add(service);
            }
        }
        return completed;
    }

    /**
     * Notify a LifecycleAware with onStart() if the system is started and the plugin publishing the service is enabled.
     *
     * @param service the service to attempt to start
     * @return true iff the service was started or is stale, false if it should be retained for later processing.
     */
    private boolean notifyOnStartIfStartedAndEnabled(final ServiceReference<LifecycleAware> service) {
        if (started) {
            return notifyLifecyleAware(service,
                    new Consumer<LifecycleAware>() {
                        @Override
                        public void accept(final LifecycleAware lifecycleAware) {
                            lifecycleAware.onStart();
                        }

                        @Override
                        public String toString() {
                            return "onStart()";
                        }
                    });
        } else {
            // We have not called onStart and need to later (when the system is started)
            return false;
        }
    }

    private boolean notifyLifecyleAware(final ServiceReference<LifecycleAware> service,
                                        final Consumer<LifecycleAware> event) {
        final Bundle bundle = service.getBundle();
        final LifecycleAware lifecycleAware = bundleContext.getService(service);
        try {
            // Either of bundle or lifecycleAware can be null if the service isn't registered.
            if ((null != bundle) && (null != lifecycleAware)) {
                final String pluginKey = getPluginKey(bundle);
                if (pluginAccessor.isPluginEnabled(pluginKey)) {
                    try {
                        log.debug("Calling LifecycleAware.{} '{}' from plugin '{}'",
                                event, lifecycleAware, pluginKey);
                        event.accept(lifecycleAware);
                    } catch (final Throwable ex) {
                        log.error("LifecycleAware.{} failed for component with class '{}' from plugin '{}'",
                                event, lifecycleAware.getClass().getName(), pluginKey, ex);
                    }
                    // We have called event
                    return true;
                } else {
                    // We have not called event and need to do it later
                    //  - for onStart(), when the plugin enables
                    //  - for onStop(), as the last resort, when service is unregistered
                    return false;
                }
            } else {
                // This can happen if we take our snapshot of services to notify, and then one is unregistered. In this case,
                // likely both bundle and lifecycleAware are null, so there's not much we can get in the message here, and it means
                // niceties like trying to safely dig data out of them is probably not worth a whole lot.
                log.warn("Discarding {} for stale LifecycleAware", event);
                // This reference is stale, so we report true as it can be discarded
                return true;
            }
        } finally {
            if (null != lifecycleAware) {
                bundleContext.ungetService(service);
            }
        }
    }

    private boolean notifyOnStopIfEnabled(final ServiceReference<LifecycleAware> service) {
        return notifyLifecyleAware(service,
                new Consumer<LifecycleAware>() {
                    @Override
                    public void accept(final LifecycleAware lifecycleAware) {
                        try {
                            // onStop() is only @since 3.0, try to call it but if it's not there then it's OK
                            lifecycleAware.onStop();
                        } catch (AbstractMethodError e) {
                            final String errorMessage = "Failed to notify with LifecycleAware.onStop()";
                            try {
                                // the method will be there because at least the LifecycleAware interface has it
                                Method onStop = lifecycleAware.getClass().getMethod("onStop");
                                if (onStop.getDeclaringClass() == LifecycleAware.class) {
                                    // expected, the message is for debug if we need it
                                    log.debug(errorMessage, e);
                                } else {
                                    // onStop() is implemented, so we need to rethrow the exception
                                    throw e;
                                }
                            } catch (NoSuchMethodException e1) {
                                // not expected
                                log.warn(errorMessage, e1);
                            }
                        }
                    }

                    @Override
                    public String toString() {
                        return "onStop()";
                    }
                });
    }

    private class LifecycleAwareServiceListener implements ServiceListener {
        /**
         * Handle a service changed by updating the list of pendingOnStart/pendingOnStop or dispatching onStart/onStop as necessary.
         *
         * @param serviceEvent details of the changed service
         */
        @Override
        public void serviceChanged(final ServiceEvent serviceEvent) {
            @SuppressWarnings("unchecked")
            // The cast here is safe since we filter by this object class when adding the service listener
            final ServiceReference<LifecycleAware> service = (ServiceReference<LifecycleAware>) serviceEvent.getServiceReference();
            switch (serviceEvent.getType()) {
                case ServiceEvent.REGISTERED: {
                    if (!notifyOnStartIfStartedAndEnabled(service)) {
                        // We didn't satisfy notification conditions, so hang on to the service for later.
                        pendingOnStart.add(service);
                    }
                    // else we tried to onStart it (we're started and the plugin is enabled), or it's stale. Either way, we're done.
                    break;
                }
                case ServiceEvent.UNREGISTERING: {
                    // It's possible we still have a reference to the service, for example if the plugin got as far as
                    // registering the service before failing, and hence never became enabled. See the comment on the
                    // pendingOnStart field for an explanation of why removing the ServiceReference directly is safe.
                    pendingOnStart.remove(service);

                    // If we failed to notify with onStop earlier, it's our last chance
                    if (pendingOnStop.remove(service)) {
                        log.warn("Notifying with LifecycleAware.onStop() on service unregister");
                        if (!notifyOnStopIfEnabled(service)) {
                            // We failed to notify onStop(), most likely due to plugin has been disabled by now
                            final Bundle bundle = service.getBundle();
                            log.warn("Failed to notify {} with LifecycleAware.onStop()", LifecycleLog.getPluginKeyFromBundle(bundle));
                        }
                    }
                    break;
                }
                default: {
                    // We filter on objectClass, which org.osgi.framework.ServiceRegistration.setProperties() says cannot be
                    // changed dynamically, so we don't need to handle MODIFIED or MODIFIED_ENDMATCH.
                    break;
                }
            }
        }
    }
}