DefaultPluginUpgradeManager

package com.atlassian.sal.core.upgrade;

import com.atlassian.beehive.ClusterLock;
import com.atlassian.beehive.ClusterLockService;
import com.atlassian.plugin.Plugin;
import com.atlassian.plugin.PluginAccessor;
import com.atlassian.plugin.event.PluginEventListener;
import com.atlassian.plugin.event.PluginEventManager;
import com.atlassian.plugin.event.events.PluginEnabledEvent;
import com.atlassian.sal.api.lifecycle.LifecycleAware;
import com.atlassian.sal.api.message.Message;
import com.atlassian.sal.api.pluginsettings.PluginSettingsFactory;
import com.atlassian.sal.api.transaction.TransactionCallback;
import com.atlassian.sal.api.transaction.TransactionTemplate;
import com.atlassian.sal.api.upgrade.PluginUpgradeManager;
import com.atlassian.sal.api.upgrade.PluginUpgradeTask;
import com.atlassian.sal.core.message.DefaultMessage;
import com.google.common.collect.ImmutableList;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.InitializingBean;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;

/**
 * Processes plugin upgrade operations.
 * <p>
 * Upgrades are triggered by the start lifecycle event, and plugin enabled.
 */
public class DefaultPluginUpgradeManager implements PluginUpgradeManager, LifecycleAware, InitializingBean {
    private static final Logger log = LoggerFactory.getLogger(DefaultPluginUpgradeManager.class);

    protected static final String LOCK_TIMEOUT_PROPERTY = "sal.upgrade.task.lock.timeout";
    protected static final int LOCK_TIMEOUT_SECONDS = Integer.getInteger(LOCK_TIMEOUT_PROPERTY, 300000);

    private volatile boolean started = false;

    private final List<PluginUpgradeTask> upgradeTasks;
    private final TransactionTemplate transactionTemplate;
    private final PluginAccessor pluginAccessor;
    private final PluginSettingsFactory pluginSettingsFactory;
    private final PluginEventManager pluginEventManager;
    private final ClusterLockService clusterLockService;

    public DefaultPluginUpgradeManager(final List<PluginUpgradeTask> upgradeTasks, final TransactionTemplate transactionTemplate,
                                       final PluginAccessor pluginAccessor, final PluginSettingsFactory pluginSettingsFactory,
                                       final PluginEventManager pluginEventManager, final ClusterLockService clusterLockService) {
        this.upgradeTasks = upgradeTasks;
        this.transactionTemplate = transactionTemplate;
        this.pluginAccessor = pluginAccessor;
        this.pluginSettingsFactory = pluginSettingsFactory;
        this.pluginEventManager = pluginEventManager;
        this.clusterLockService = clusterLockService;
    }

    /**
     * Notifies the plugin upgrade manager that a plugin update tasks has been registered.
     * <p>
     * This method does nothing but logging at the moment. Is is now deprecated since it could result in circular
     * dependency when trying to bind already exposed update tasks to plugin manager that is being created.
     *
     * @param task  the upgrade task that is being bound
     * @param props the set of properties that the upgrade task was registered with
     * @deprecated as of 2.0.16 no longer used we now use the eventing framework to run upgrade tasks for plugins that
     *             are 'enabled', see {@link #onPluginEnabled(com.atlassian.plugin.event.events.PluginEnabledEvent)}
     */
    @SuppressWarnings("unchecked")
    @Deprecated
    public void onBind(final PluginUpgradeTask task, final Map props) {
        // Doing lots here....
        log.debug("onbind task = [" + task.getPluginKey() + ", " + task.getBuildNumber() + "] ");
    }

    @Override
    public void onStart() {
        log.debug("onStart");
        final List<Message> messages = upgrade();

        // TODO 1: should do something useful with the messages
        // TODO 2: we don't know what upgrade tasks these messages came from
        if (messages != null) {
            for (final Message msg : messages) {
                log.error("Upgrade error: " + msg);
            }
        }

        started = true;
    }

    @Override
    public void onStop() {
        pluginEventManager.unregister(this);
    }

    @PluginEventListener
    public void onPluginEnabled(PluginEnabledEvent event) {
        // Check if the Application is fully started:
        if (started) {
            // Run upgrades for this plugin that as been enabled AFTER the onStart event.
            final List<Message> messages = upgradeInternal(event.getPlugin());
            if (messages != null && messages.size() > 0) {
                log.error("Error(s) encountered while upgrading plugin '" + event.getPlugin().getName() + "' on enable.");
                for (final Message msg : messages) {
                    log.error("Upgrade error: " + msg);
                }
            }
        }
        // If onStart() has not occurred yet then ignore event - we need to wait until the App is started properly.
    }

    /**
     * @return map of all upgrade tasks (stored by pluginKey)
     */
    protected Map<String, List<PluginUpgradeTask>> getUpgradeTasks() {
        final Map<String, List<PluginUpgradeTask>> pluginUpgrades = new HashMap<>();

        // Find all implementations of PluginUpgradeTask
        for (final PluginUpgradeTask upgradeTask : upgradeTasks) {
            pluginUpgrades.computeIfAbsent(upgradeTask.getPluginKey(), k -> new ArrayList<>())
                    .add(upgradeTask);
        }

        return pluginUpgrades;
    }


    @SuppressWarnings("unchecked")
    public List<Message> upgrade() {
        return upgradeInternal();
    }

    public List<Message> upgradeInternal() {
        log.info("Running plugin upgrade tasks...");

        // 1. get all upgrade tasks for all plugins
        final Map<String, List<PluginUpgradeTask>> pluginUpgrades = getUpgradeTasks();


        final ArrayList<Message> messages = new ArrayList<Message>();

        // 2. for each plugin, sort tasks by build number and execute them
        for (final String pluginKey : pluginUpgrades.keySet()) {
            final List<PluginUpgradeTask> upgrades = pluginUpgrades.get(pluginKey);

            final List<Message> upgradeMessages = upgradePlugin(pluginKey, upgrades);
            if (upgradeMessages != null) {
                messages.addAll(upgradeMessages);
            }
        }

        return messages;
    }

    public List<Message> upgradeInternal(Plugin plugin) {
        final Map<String, List<PluginUpgradeTask>> pluginUpgrades = getUpgradeTasks();
        final String pluginKey = plugin.getKey();
        final List<PluginUpgradeTask> upgrades = pluginUpgrades.get(pluginKey);
        if (upgrades == null) {
            // nothing to do
            return null;
        }
        return upgradePlugin(pluginKey, upgrades);
    }

    private List<Message> upgradePlugin(final String pluginKey, final List<PluginUpgradeTask> upgrades) {
        return transactionTemplate.execute(() -> {
            final Plugin plugin = pluginAccessor.getPlugin(pluginKey);
            if (plugin == null) {
                throw new IllegalArgumentException("Invalid plugin key: " + pluginKey);
            }

            final PluginUpgrader pluginUpgrader = new PluginUpgrader(plugin, pluginSettingsFactory.createGlobalSettings(), upgrades);

            final String lockName = "sal.upgrade." + pluginKey;
            final ClusterLock lock = clusterLockService.getLockForName(lockName);
            try {
                if (!lock.tryLock(LOCK_TIMEOUT_SECONDS, TimeUnit.SECONDS)) {
                    final String timeoutMessage = "unable to acquire cluster lock named '" + lockName +
                            "' after waiting " + LOCK_TIMEOUT_SECONDS +
                            " seconds; note that this timeout may be adjusted via the system property '" +
                            LOCK_TIMEOUT_PROPERTY + "'";
                    log.error(timeoutMessage);
                    return ImmutableList.of(new DefaultMessage(timeoutMessage));
                }
            } catch (InterruptedException e) {
                final String interruptedMessage = "interrupted while trying to acquire cluster lock named '" + lockName +
                        "' " + e.getMessage();
                log.error(interruptedMessage);
                return ImmutableList.of(new DefaultMessage(interruptedMessage));
            }

            try {
                return pluginUpgrader.upgrade();
            } finally {
                lock.unlock();
            }
        });
    }

    public void afterPropertiesSet() throws Exception {
        pluginEventManager.register(this);
    }
}