Merge "Provide mechanism for on demand loading/unloading of connected plugins" into udc-dev

packages/SystemUI/customization/src/com/android/systemui/shared/clocks/ClockRegistry.kt

@@ -29,14 +29,15 @@ import com.android.systemui.plugins.ClockMetadata
 import com.android.systemui.plugins.ClockProvider
 import com.android.systemui.plugins.ClockProviderPlugin
 import com.android.systemui.plugins.ClockSettings
+import com.android.systemui.plugins.PluginLifecycleManager
 import com.android.systemui.plugins.PluginListener
 import com.android.systemui.plugins.PluginManager
 import com.android.systemui.util.Assert
+import java.util.concurrent.atomic.AtomicBoolean
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.launch
 
-private val TAG = ClockRegistry::class.simpleName!!
 private const val DEBUG = true
 private val KEY_TIMESTAMP = "appliedTimestamp"
 
@@ -51,7 +52,10 @@ open class ClockRegistry(
     val handleAllUsers: Boolean,
     defaultClockProvider: ClockProvider,
     val fallbackClockId: ClockId = DEFAULT_CLOCK_ID,
+    val keepAllLoaded: Boolean,
+    val subTag: String,
 ) {
+    private val TAG = "${ClockRegistry::class.simpleName} ($subTag)"
     interface ClockChangeListener {
         // Called when the active clock changes
         fun onCurrentClockChanged() {}
@@ -76,11 +80,85 @@ open class ClockRegistry(
 
     private val pluginListener =
         object : PluginListener<ClockProviderPlugin> {
-            override fun onPluginConnected(plugin: ClockProviderPlugin, context: Context) =
+            override fun onPluginAttached(manager: PluginLifecycleManager<ClockProviderPlugin>) {
-                connectClocks(plugin)
+                manager.loadPlugin()
+            }
 
-            override fun onPluginDisconnected(plugin: ClockProviderPlugin) =
+            override fun onPluginLoaded(
-                disconnectClocks(plugin)
+                plugin: ClockProviderPlugin,
+                pluginContext: Context,
+                manager: PluginLifecycleManager<ClockProviderPlugin>
+            ) {
+                var isClockListChanged = false
+                for (clock in plugin.getClocks()) {
+                    val id = clock.clockId
+                    var isNew = false
+                    val info =
+                        availableClocks.getOrPut(id) {
+                            isNew = true
+                            ClockInfo(clock, plugin, manager)
+                        }
+
+                    if (isNew) {
+                        isClockListChanged = true
+                        onConnected(id)
+                    }
+
+                    if (manager != info.manager) {
+                        Log.e(
+                            TAG,
+                            "Clock Id conflict on load: $id is registered to another provider"
+                        )
+                        continue
+                    }
+
+                    info.provider = plugin
+                    onLoaded(id)
+                }
+
+                if (isClockListChanged) {
+                    triggerOnAvailableClocksChanged()
+                }
+                verifyLoadedProviders()
+            }
+
+            override fun onPluginUnloaded(
+                plugin: ClockProviderPlugin,
+                manager: PluginLifecycleManager<ClockProviderPlugin>
+            ) {
+                for (clock in plugin.getClocks()) {
+                    val id = clock.clockId
+                    val info = availableClocks[id]
+                    if (info?.manager != manager) {
+                        Log.e(
+                            TAG,
+                            "Clock Id conflict on unload: $id is registered to another provider"
+                        )
+                        continue
+                    }
+                    info.provider = null
+                    onUnloaded(id)
+                }
+
+                verifyLoadedProviders()
+            }
+
+            override fun onPluginDetached(manager: PluginLifecycleManager<ClockProviderPlugin>) {
+                val removed = mutableListOf<ClockId>()
+                availableClocks.entries.removeAll {
+                    if (it.value.manager != manager) {
+                        return@removeAll false
+                    }
+
+                    removed.add(it.key)
+                    return@removeAll true
+                }
+
+                removed.forEach(::onDisconnected)
+                if (removed.size > 0) {
+                    triggerOnAvailableClocksChanged()
+                }
+            }
         }
 
     private val userSwitchObserver =
@@ -96,7 +174,8 @@ open class ClockRegistry(
         protected set(value) {
             if (field != value) {
                 field = value
-                scope.launch(mainDispatcher) { onClockChanged { it.onCurrentClockChanged() } }
+                verifyLoadedProviders()
+                triggerOnCurrentClockChanged()
             }
         }
 
@@ -168,9 +247,36 @@ open class ClockRegistry(
         Assert.isNotMainThread()
     }
 
-    private fun onClockChanged(func: (ClockChangeListener) -> Unit) {
+    private var isClockChanged = AtomicBoolean(false)
-        assertMainThread()
+    private fun triggerOnCurrentClockChanged() {
-        clockChangeListeners.forEach(func)
+        val shouldSchedule = isClockChanged.compareAndSet(false, true)
+        if (!shouldSchedule) {
+            return
+        }
+
+        android.util.Log.e("HAWK", "triggerOnCurrentClockChanged")
+        scope.launch(mainDispatcher) {
+            assertMainThread()
+            android.util.Log.e("HAWK", "isClockChanged")
+            isClockChanged.set(false)
+            clockChangeListeners.forEach { it.onCurrentClockChanged() }
+        }
+    }
+
+    private var isClockListChanged = AtomicBoolean(false)
+    private fun triggerOnAvailableClocksChanged() {
+        val shouldSchedule = isClockListChanged.compareAndSet(false, true)
+        if (!shouldSchedule) {
+            return
+        }
+
+        android.util.Log.e("HAWK", "triggerOnAvailableClocksChanged")
+        scope.launch(mainDispatcher) {
+            assertMainThread()
+            android.util.Log.e("HAWK", "isClockListChanged")
+            isClockListChanged.set(false)
+            clockChangeListeners.forEach { it.onAvailableClocksChanged() }
+        }
     }
 
     public fun mutateSetting(mutator: (ClockSettings) -> ClockSettings) {
@@ -190,7 +296,12 @@ open class ClockRegistry(
         }
 
     init {
-        connectClocks(defaultClockProvider)
+        // Register default clock designs
+        for (clock in defaultClockProvider.getClocks()) {
+            availableClocks[clock.clockId] = ClockInfo(clock, defaultClockProvider, null)
+        }
+
+        // Something has gone terribly wrong if the default clock isn't present
         if (!availableClocks.containsKey(DEFAULT_CLOCK_ID)) {
             throw IllegalArgumentException(
                 "$defaultClockProvider did not register clock at $DEFAULT_CLOCK_ID"
@@ -244,59 +355,87 @@ open class ClockRegistry(
         }
     }
 
-    private fun connectClocks(provider: ClockProvider) {
+    private var isVerifying = AtomicBoolean(false)
-        var isAvailableChanged = false
+    private fun verifyLoadedProviders() {
-        val currentId = currentClockId
+        val shouldSchedule = isVerifying.compareAndSet(false, true)
-        for (clock in provider.getClocks()) {
+        if (!shouldSchedule) {
-            val id = clock.clockId
+            return
-            val current = availableClocks[id]
-            if (current != null) {
-                Log.e(
-                    TAG,
-                    "Clock Id conflict: $id is registered by both " +
-                        "${provider::class.simpleName} and ${current.provider::class.simpleName}"
-                )
-                continue
-            }
-
-            availableClocks[id] = ClockInfo(clock, provider)
-            isAvailableChanged = true
-            if (DEBUG) {
-                Log.i(TAG, "Added ${clock.clockId}")
-            }
-
-            if (currentId == id) {
-                if (DEBUG) {
-                    Log.i(TAG, "Current clock ($currentId) was connected")
-                }
-                onClockChanged { it.onCurrentClockChanged() }
-            }
         }
 
-        if (isAvailableChanged) {
+        scope.launch(bgDispatcher) {
-            onClockChanged { it.onAvailableClocksChanged() }
+            if (keepAllLoaded) {
+                // Enforce that all plugins are loaded if requested
+                for ((_, info) in availableClocks) {
+                    info.manager?.loadPlugin()
+                }
+                isVerifying.set(false)
+                return@launch
+            }
+
+            val currentClock = availableClocks[currentClockId]
+            if (currentClock == null) {
+                // Current Clock missing, load no plugins and use default
+                for ((_, info) in availableClocks) {
+                    info.manager?.unloadPlugin()
+                }
+                isVerifying.set(false)
+                return@launch
+            }
+
+            val currentManager = currentClock.manager
+            currentManager?.loadPlugin()
+
+            for ((_, info) in availableClocks) {
+                val manager = info.manager
+                if (manager != null && manager.isLoaded && currentManager != manager) {
+                    manager.unloadPlugin()
+                }
+            }
+            isVerifying.set(false)
         }
     }
 
-    private fun disconnectClocks(provider: ClockProvider) {
+    private fun onConnected(clockId: ClockId) {
-        var isAvailableChanged = false
+        if (DEBUG) {
-        val currentId = currentClockId
+            Log.i(TAG, "Connected $clockId")
-        for (clock in provider.getClocks()) {
-            availableClocks.remove(clock.clockId)
-            isAvailableChanged = true
-
-            if (DEBUG) {
-                Log.i(TAG, "Removed ${clock.clockId}")
-            }
-
-            if (currentId == clock.clockId) {
-                Log.w(TAG, "Current clock ($currentId) was disconnected")
-                onClockChanged { it.onCurrentClockChanged() }
-            }
         }
 
-        if (isAvailableChanged) {
+        if (currentClockId == clockId) {
-            onClockChanged { it.onAvailableClocksChanged() }
+            if (DEBUG) {
+                Log.i(TAG, "Current clock ($clockId) was connected")
+            }
+        }
+    }
+
+    private fun onLoaded(clockId: ClockId) {
+        if (DEBUG) {
+            Log.i(TAG, "Loaded $clockId")
+        }
+
+        if (currentClockId == clockId) {
+            Log.i(TAG, "Current clock ($clockId) was loaded")
+            triggerOnCurrentClockChanged()
+        }
+    }
+
+    private fun onUnloaded(clockId: ClockId) {
+        if (DEBUG) {
+            Log.i(TAG, "Unloaded $clockId")
+        }
+
+        if (currentClockId == clockId) {
+            Log.w(TAG, "Current clock ($clockId) was unloaded")
+            triggerOnCurrentClockChanged()
+        }
+    }
+
+    private fun onDisconnected(clockId: ClockId) {
+        if (DEBUG) {
+            Log.i(TAG, "Disconnected $clockId")
+        }
+
+        if (currentClockId == clockId) {
+            Log.w(TAG, "Current clock ($clockId) was disconnected")
         }
     }
 
@@ -345,6 +484,7 @@ open class ClockRegistry(
 
     private data class ClockInfo(
         val metadata: ClockMetadata,
-        val provider: ClockProvider,
+        var provider: ClockProvider?,
+        val manager: PluginLifecycleManager<ClockProviderPlugin>?,
     )
 }

packages/SystemUI/plugin_core/src/com/android/systemui/plugins/PluginLifecycleManager.java

@@ -0,0 +1,48 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * 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 com.android.systemui.plugins;
+
+/**
+ * Provides the ability for consumers to control plugin lifecycle.
+ *
+ * @param <T> is the target plugin type
+ */
+public interface PluginLifecycleManager<T extends Plugin> {
+    /** Returns the currently loaded plugin instance (if plugin is loaded) */
+    T getPlugin();
+
+    /** returns true if the plugin is currently loaded */
+    default boolean isLoaded() {
+        return getPlugin() != null;
+    }
+
+    /**
+     * Loads and creates the plugin instance if it does not exist.
+     *
+     * This will trigger {@link PluginListener#onPluginLoaded} with the new instance if it did not
+     * already exist.
+     */
+    void loadPlugin();
+
+    /**
+     * Unloads and destroys the plugin instance if it exists.
+     *
+     * This will trigger {@link PluginListener#onPluginUnloaded} if a concrete plugin instance
+     * existed when this call was made.
+     */
+    void unloadPlugin();
+}

packages/SystemUI/plugin_core/src/com/android/systemui/plugins/PluginListener.java

@@ -17,7 +17,32 @@ package com.android.systemui.plugins;
 import android.content.Context;
 
 /**
- * Interface for listening to plugins being connected.
+ * Interface for listening to plugins being connected and disconnected.
+ *
+ * The call order for a plugin is
+ *  1) {@link #onPluginAttached}
+ *          Called when a new plugin is added to the device, or an existing plugin was replaced by
+ *          the package manager. Will only be called once per package manager event. If multiple
+ *          non-conflicting packages which have the same plugin interface are installed on the
+ *          device, then this method can be called multiple times with different instances of
+ *          {@link PluginLifecycleManager} (as long as `allowMultiple` was set to true when the
+ *          listener was registered with {@link PluginManager#addPluginListener}).
+ *  2) {@link #onPluginLoaded}
+ *          Called whenever a new instance of the plugin object is created and ready for use. Can be
+ *          called multiple times per {@link PluginLifecycleManager}, but will always pass a newly
+ *          created plugin object. {@link #onPluginUnloaded} with the previous plugin object will
+ *          be called before another call to {@link #onPluginLoaded} is made. This method will be
+ *          called once automatically after {@link #onPluginAttached}. Besides the initial call,
+ *          {@link #onPluginLoaded} will occur due to {@link PluginLifecycleManager#loadPlugin}.
+ *  3) {@link #onPluginUnloaded}
+ *          Called when a request to unload the plugin has been received. This can be triggered from
+ *          a related call to {@link PluginLifecycleManager#unloadPlugin} or for any reason that
+ *          {@link #onPluginDetached} would be triggered.
+ *  4) {@link #onPluginDetached}
+ *          Called when the package is removed from the device, disabled, or replaced due to an
+ *          external trigger. These are events from the android package manager.
+ *
+ * @param <T> is the target plugin type
  */
 public interface PluginListener<T extends Plugin> {
     /**
@@ -25,14 +50,69 @@ public interface PluginListener<T extends Plugin> {
      * This may be called multiple times if multiple plugins are allowed.
      * It may also be called in the future if the plugin package changes
      * and needs to be reloaded.
+     *
+     * @deprecated Migrate to {@link #onPluginLoaded} or {@link #onPluginAttached}
      */
-    void onPluginConnected(T plugin, Context pluginContext);
+    @Deprecated
+    default void onPluginConnected(T plugin, Context pluginContext) {
+        // Optional
+    }
+
+    /**
+     * Called when the plugin is first attached to the host application. {@link #onPluginLoaded}
+     * will be automatically called as well when first attached. This may be called multiple times
+     * if multiple plugins are allowed. It may also be called in the future if the plugin package
+     * changes and needs to be reloaded. Each call to {@link #onPluginAttached} will provide a new
+     * or different {@link PluginLifecycleManager}.
+     */
+    default void onPluginAttached(PluginLifecycleManager<T> manager) {
+        // Optional
+    }
 
     /**
      * Called when a plugin has been uninstalled/updated and should be removed
      * from use.
+     *
+     * @deprecated Migrate to {@link #onPluginDetached} or {@link #onPluginUnloaded}
      */
+    @Deprecated
     default void onPluginDisconnected(T plugin) {
         // Optional.
     }
-}
+
+    /**
+     * Called when the plugin has been detached from the host application. Implementers should no
+     * longer attempt to reload it via this {@link PluginLifecycleManager}. If the package was
+     * updated and not removed, then {@link #onPluginAttached} will be called again when the updated
+     * package is available.
+     */
+    default void onPluginDetached(PluginLifecycleManager<T> manager) {
+        // Optional.
+    }
+
+    /**
+     * Called when the plugin is loaded into the host's process and is available for use. This can
+     * happen several times if clients are using {@link PluginLifecycleManager} to manipulate a
+     * plugin's load state. Each call to {@link #onPluginLoaded} will have a matched call to
+     * {@link #onPluginUnloaded} when that plugin object should no longer be used.
+     */
+    default void onPluginLoaded(
+            T plugin,
+            Context pluginContext,
+            PluginLifecycleManager<T> manager
+    ) {
+        // Optional, default to deprecated version
+        onPluginConnected(plugin, pluginContext);
+    }
+
+    /**
+     * Called when the plugin should no longer be used. Listeners should clean up all references to
+     * the relevant plugin so that it can be garbage collected. If the plugin object is required in
+     * the future a call can be made to {@link PluginLifecycleManager#loadPlugin} to create a new
+     * plugin object and trigger {@link #onPluginLoaded}.
+     */
+    default void onPluginUnloaded(T plugin, PluginLifecycleManager<T> manager) {
+        // Optional, default to deprecated version
+        onPluginDisconnected(plugin);
+    }
+}

packages/SystemUI/shared/src/com/android/systemui/shared/plugins/PluginActionManager.java

@@ -210,12 +210,12 @@ public class PluginActionManager<T extends Plugin> {
     private void onPluginConnected(PluginInstance<T> pluginInstance) {
         if (DEBUG) Log.d(TAG, "onPluginConnected");
         PluginPrefs.setHasPlugins(mContext);
-        pluginInstance.onCreate(mContext, mListener);
+        pluginInstance.onCreate();
     }
 
     private void onPluginDisconnected(PluginInstance<T> pluginInstance) {
         if (DEBUG) Log.d(TAG, "onPluginDisconnected");
-        pluginInstance.onDestroy(mListener);
+        pluginInstance.onDestroy();
     }
 
     private void queryAll() {
@@ -312,7 +312,7 @@ public class PluginActionManager<T extends Plugin> {
             try {
                 return mPluginInstanceFactory.create(
                         mContext, appInfo, component,
-                        mPluginClass);
+                        mPluginClass, mListener);
             } catch (InvalidVersionException e) {
                 reportInvalidVersion(component, component.getClassName(), e);
             }

packages/SystemUI/shared/src/com/android/systemui/shared/plugins/PluginInstance.java

@@ -21,13 +21,16 @@ import android.content.ComponentName;
 import android.content.Context;
 import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageManager;
+import android.content.pm.PackageManager.NameNotFoundException;
 import android.text.TextUtils;
-import android.util.ArrayMap;
 import android.util.Log;
 
+import androidx.annotation.Nullable;
+
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.systemui.plugins.Plugin;
 import com.android.systemui.plugins.PluginFragment;
+import com.android.systemui.plugins.PluginLifecycleManager;
 import com.android.systemui.plugins.PluginListener;
 
 import dalvik.system.PathClassLoader;
@@ -35,7 +38,7 @@ import dalvik.system.PathClassLoader;
 import java.io.File;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.Map;
+import java.util.function.Supplier;
 
 /**
  * Contains a single instantiation of a Plugin.
@@ -45,42 +48,102 @@ import java.util.Map;
  *
  * @param <T> The type of plugin that this contains.
  */
-public class PluginInstance<T extends Plugin> {
+public class PluginInstance<T extends Plugin> implements PluginLifecycleManager {
     private static final String TAG = "PluginInstance";
-    private static final Map<String, ClassLoader> sClassLoaders = new ArrayMap<>();
 
-    private final Context mPluginContext;
+    private final Context mAppContext;
-    private final VersionInfo mVersionInfo;
+    private final PluginListener<T> mListener;
     private final ComponentName mComponentName;
-    private final T mPlugin;
+    private final PluginFactory<T> mPluginFactory;
+
+    private Context mPluginContext;
+    private T mPlugin;
 
     /** */
-    public PluginInstance(ComponentName componentName, T plugin, Context pluginContext,
+    public PluginInstance(
-            VersionInfo versionInfo) {
+            Context appContext,
+            PluginListener<T> listener,
+            ComponentName componentName,
+            PluginFactory<T> pluginFactory,
+            @Nullable T plugin) {
+        mAppContext = appContext;
+        mListener = listener;
         mComponentName = componentName;
+        mPluginFactory = pluginFactory;
         mPlugin = plugin;
-        mPluginContext = pluginContext;
+
-        mVersionInfo = versionInfo;
+        if (mPlugin != null) {
+            mPluginContext = mPluginFactory.createPluginContext();
+        }
     }
 
     /** Alerts listener and plugin that the plugin has been created. */
-    public void onCreate(Context appContext, PluginListener<T> listener) {
+    public void onCreate() {
-        if (!(mPlugin instanceof PluginFragment)) {
+        mListener.onPluginAttached(this);
-            // Only call onCreate for plugins that aren't fragments, as fragments
+        if (mPlugin == null) {
-            // will get the onCreate as part of the fragment lifecycle.
+            loadPlugin();
-            mPlugin.onCreate(appContext, mPluginContext);
+        } else {
+            if (!(mPlugin instanceof PluginFragment)) {
+                // Only call onCreate for plugins that aren't fragments, as fragments
+                // will get the onCreate as part of the fragment lifecycle.
+                mPlugin.onCreate(mAppContext, mPluginContext);
+            }
+            mListener.onPluginLoaded(mPlugin, mPluginContext, this);
         }
-        listener.onPluginConnected(mPlugin, mPluginContext);
     }
 
     /** Alerts listener and plugin that the plugin is being shutdown. */
-    public void onDestroy(PluginListener<T> listener) {
+    public void onDestroy() {
-        listener.onPluginDisconnected(mPlugin);
+        unloadPlugin();
+        mListener.onPluginDetached(this);
+    }
+
+    /** Returns the current plugin instance (if it is loaded). */
+    @Nullable
+    public T getPlugin() {
+        return mPlugin;
+    }
+
+    /**
+     * Loads and creates the plugin if it does not exist.
+     */
+    public void loadPlugin() {
+        if (mPlugin != null) {
+            return;
+        }
+
+        mPlugin = mPluginFactory.createPlugin();
+        mPluginContext = mPluginFactory.createPluginContext();
+        if (mPlugin == null || mPluginContext == null) {
+            return;
+        }
+
+        if (!(mPlugin instanceof PluginFragment)) {
+            // Only call onCreate for plugins that aren't fragments, as fragments
+            // will get the onCreate as part of the fragment lifecycle.
+            mPlugin.onCreate(mAppContext, mPluginContext);
+        }
+        mListener.onPluginLoaded(mPlugin, mPluginContext, this);
+    }
+
+    /**
+     * Unloads and destroys the current plugin instance if it exists.
+     *
+     * This will free the associated memory if there are not other references.
+     */
+    public void unloadPlugin() {
+        if (mPlugin == null) {
+            return;
+        }
+
+        mListener.onPluginUnloaded(mPlugin, this);
         if (!(mPlugin instanceof PluginFragment)) {
             // Only call onDestroy for plugins that aren't fragments, as fragments
             // will get the onDestroy as part of the fragment lifecycle.
             mPlugin.onDestroy();
         }
+        mPlugin = null;
+        mPluginContext = null;
     }
 
     /**
@@ -89,7 +152,7 @@ public class PluginInstance<T extends Plugin> {
      * It does this by string comparison of the class names.
      **/
     public boolean containsPluginClass(Class pluginClass) {
-        return mPlugin.getClass().getName().equals(pluginClass.getName());
+        return mComponentName.getClassName().equals(pluginClass.getName());
     }
 
     public ComponentName getComponentName() {
@@ -101,7 +164,7 @@ public class PluginInstance<T extends Plugin> {
     }
 
     public VersionInfo getVersionInfo() {
-        return mVersionInfo;
+        return mPluginFactory.checkVersion(mPlugin);
     }
 
     @VisibleForTesting
@@ -134,21 +197,20 @@ public class PluginInstance<T extends Plugin> {
                 Context context,
                 ApplicationInfo appInfo,
                 ComponentName componentName,
-                Class<T> pluginClass)
+                Class<T> pluginClass,
+                PluginListener<T> listener)
                 throws PackageManager.NameNotFoundException, ClassNotFoundException,
                 InstantiationException, IllegalAccessException {
 
-            ClassLoader classLoader = getClassLoader(appInfo, mBaseClassLoader);
+            PluginFactory<T> pluginFactory = new PluginFactory<T>(
-            Context pluginContext = new PluginActionManager.PluginContextWrapper(
+                    context, mInstanceFactory, appInfo, componentName, mVersionChecker, pluginClass,
-                    context.createApplicationContext(appInfo, 0), classLoader);
+                    () -> getClassLoader(appInfo, mBaseClassLoader));
-            Class<T> instanceClass = (Class<T>) Class.forName(
-                    componentName.getClassName(), true, classLoader);
             // TODO: Only create the plugin before version check if we need it for
             // legacy version check.
-            T instance = (T) mInstanceFactory.create(instanceClass);
+            T instance = pluginFactory.createPlugin();
-            VersionInfo version = mVersionChecker.checkVersion(
+            pluginFactory.checkVersion(instance);
-                    instanceClass, pluginClass, instance);
+            return new PluginInstance<T>(
-            return new PluginInstance<T>(componentName, instance, pluginContext, version);
+                    context, listener, componentName, pluginFactory, instance);
         }
 
         private boolean isPluginPackagePrivileged(String packageName) {
@@ -179,9 +241,6 @@ public class PluginInstance<T extends Plugin> {
                         + appInfo.sourceDir + ", pkg: " + appInfo.packageName);
                 return null;
             }
-            if (sClassLoaders.containsKey(appInfo.packageName)) {
-                return sClassLoaders.get(appInfo.packageName);
-            }
 
             List<String> zipPaths = new ArrayList<>();
             List<String> libPaths = new ArrayList<>();
@@ -190,13 +249,20 @@ public class PluginInstance<T extends Plugin> {
                     TextUtils.join(File.pathSeparator, zipPaths),
                     TextUtils.join(File.pathSeparator, libPaths),
                     getParentClassLoader(baseClassLoader));
-            sClassLoaders.put(appInfo.packageName, classLoader);
             return classLoader;
         }
     }
 
     /** Class that compares a plugin class against an implementation for version matching. */
-    public static class VersionChecker {
+    public interface VersionChecker {
+        /** Compares two plugin classes. */
+        <T extends Plugin> VersionInfo checkVersion(
+                Class<T> instanceClass, Class<T> pluginClass, Plugin plugin);
+    }
+
+    /** Class that compares a plugin class against an implementation for version matching. */
+    public static class VersionCheckerImpl implements VersionChecker {
+        @Override
         /** Compares two plugin classes. */
         public <T extends Plugin> VersionInfo checkVersion(
                 Class<T> instanceClass, Class<T> pluginClass, Plugin plugin) {
@@ -204,7 +270,7 @@ public class PluginInstance<T extends Plugin> {
             VersionInfo instanceVersion = new VersionInfo().addClass(instanceClass);
             if (instanceVersion.hasVersionInfo()) {
                 pluginVersion.checkVersion(instanceVersion);
-            } else {
+            } else if (plugin != null) {
                 int fallbackVersion = plugin.getVersion();
                 if (fallbackVersion != pluginVersion.getDefaultVersion()) {
                     throw new VersionInfo.InvalidVersionException("Invalid legacy version", false);
@@ -225,4 +291,74 @@ public class PluginInstance<T extends Plugin> {
             return (T) cls.newInstance();
         }
     }
+
+    /**
+     * Instanced wrapper of InstanceFactory
+     *
+     * @param <T> is the type of the plugin object to be built
+     **/
+    public static class PluginFactory<T extends Plugin> {
+        private final Context mContext;
+        private final InstanceFactory<?> mInstanceFactory;
+        private final ApplicationInfo mAppInfo;
+        private final ComponentName mComponentName;
+        private final VersionChecker mVersionChecker;
+        private final Class<T> mPluginClass;
+        private final Supplier<ClassLoader> mClassLoaderFactory;
+
+        public PluginFactory(
+                Context context,
+                InstanceFactory<?> instanceFactory,
+                ApplicationInfo appInfo,
+                ComponentName componentName,
+                VersionChecker versionChecker,
+                Class<T> pluginClass,
+                Supplier<ClassLoader> classLoaderFactory) {
+            mContext = context;
+            mInstanceFactory = instanceFactory;
+            mAppInfo = appInfo;
+            mComponentName = componentName;
+            mVersionChecker = versionChecker;
+            mPluginClass = pluginClass;
+            mClassLoaderFactory = classLoaderFactory;
+        }
+
+        /** Creates the related plugin object from the factory */
+        public T createPlugin() {
+            try {
+                ClassLoader loader = mClassLoaderFactory.get();
+                Class<T> instanceClass = (Class<T>) Class.forName(
+                        mComponentName.getClassName(), true, loader);
+                return (T) mInstanceFactory.create(instanceClass);
+            } catch (ClassNotFoundException ex) {
+                Log.e(TAG, "Failed to load plugin", ex);
+            } catch (IllegalAccessException ex) {
+                Log.e(TAG, "Failed to load plugin", ex);
+            } catch (InstantiationException ex) {
+                Log.e(TAG, "Failed to load plugin", ex);
+            }
+            return null;
+        }
+
+        /** Creates a context wrapper for the plugin */
+        public Context createPluginContext() {
+            try {
+                ClassLoader loader = mClassLoaderFactory.get();
+                return new PluginActionManager.PluginContextWrapper(
+                    mContext.createApplicationContext(mAppInfo, 0), loader);
+            } catch (NameNotFoundException ex) {
+                Log.e(TAG, "Failed to create plugin context", ex);
+            }
+            return null;
+        }
+
+        /** Check Version and create VersionInfo for instance */
+        public VersionInfo checkVersion(T instance) {
+            if (instance == null) {
+                instance = createPlugin();
+            }
+            return mVersionChecker.checkVersion(
+                    (Class<T>) instance.getClass(), mPluginClass, instance);
+        }
+    }
 }

packages/SystemUI/src/com/android/keyguard/dagger/ClockRegistryModule.java

@@ -60,7 +60,9 @@ public abstract class ClockRegistryModule {
                 featureFlags.isEnabled(Flags.LOCKSCREEN_CUSTOM_CLOCKS),
                 /* handleAllUsers= */ true,
                 new DefaultClockProvider(context, layoutInflater, resources),
-                context.getString(R.string.lockscreen_clock_id_fallback));
+                context.getString(R.string.lockscreen_clock_id_fallback),
+                /* keepAllLoaded = */ false,
+                /* subTag = */ "System");
         registry.registerListeners();
         return registry;
     }

packages/SystemUI/src/com/android/systemui/plugins/PluginsModule.java

@@ -73,7 +73,7 @@ public abstract class PluginsModule {
         return new PluginInstance.Factory(
                 PluginModule.class.getClassLoader(),
                 new PluginInstance.InstanceFactory<>(),
-                new PluginInstance.VersionChecker(),
+                new PluginInstance.VersionCheckerImpl(),
                 privilegedPlugins,
                 isDebug);
     }

packages/SystemUI/tests/src/com/android/systemui/shared/clocks/ClockRegistryTest.kt

@@ -27,13 +27,16 @@ import com.android.systemui.plugins.ClockMetadata
 import com.android.systemui.plugins.ClockProviderPlugin
 import com.android.systemui.plugins.ClockSettings
 import com.android.systemui.plugins.PluginListener
+import com.android.systemui.plugins.PluginLifecycleManager
 import com.android.systemui.plugins.PluginManager
 import com.android.systemui.util.mockito.argumentCaptor
 import com.android.systemui.util.mockito.eq
+import com.android.systemui.util.mockito.mock
 import junit.framework.Assert.assertEquals
 import junit.framework.Assert.fail
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.test.TestCoroutineScheduler
 import kotlinx.coroutines.test.TestScope
 import org.junit.Before
 import org.junit.Rule
@@ -49,6 +52,7 @@ import org.mockito.junit.MockitoJUnit
 class ClockRegistryTest : SysuiTestCase() {
 
     @JvmField @Rule val mockito = MockitoJUnit.rule()
+    private lateinit var scheduler: TestCoroutineScheduler
     private lateinit var dispatcher: CoroutineDispatcher
     private lateinit var scope: TestScope
 
@@ -58,37 +62,38 @@ class ClockRegistryTest : SysuiTestCase() {
     @Mock private lateinit var mockDefaultClock: ClockController
     @Mock private lateinit var mockThumbnail: Drawable
     @Mock private lateinit var mockContentResolver: ContentResolver
+    @Mock private lateinit var mockPluginLifecycle: PluginLifecycleManager<ClockProviderPlugin>
     private lateinit var fakeDefaultProvider: FakeClockPlugin
     private lateinit var pluginListener: PluginListener<ClockProviderPlugin>
     private lateinit var registry: ClockRegistry
 
     companion object {
-        private fun failFactory(): ClockController {
+        private fun failFactory(clockId: ClockId): ClockController {
-            fail("Unexpected call to createClock")
+            fail("Unexpected call to createClock: $clockId")
             return null!!
         }
 
-        private fun failThumbnail(): Drawable? {
+        private fun failThumbnail(clockId: ClockId): Drawable? {
-            fail("Unexpected call to getThumbnail")
+            fail("Unexpected call to getThumbnail: $clockId")
             return null
         }
     }
 
     private class FakeClockPlugin : ClockProviderPlugin {
         private val metadata = mutableListOf<ClockMetadata>()
-        private val createCallbacks = mutableMapOf<ClockId, () -> ClockController>()
+        private val createCallbacks = mutableMapOf<ClockId, (ClockId) -> ClockController>()
-        private val thumbnailCallbacks = mutableMapOf<ClockId, () -> Drawable?>()
+        private val thumbnailCallbacks = mutableMapOf<ClockId, (ClockId) -> Drawable?>()
 
         override fun getClocks() = metadata
         override fun createClock(settings: ClockSettings): ClockController =
-            createCallbacks[settings.clockId!!]!!()
+            createCallbacks[settings.clockId!!]!!(settings.clockId!!)
-        override fun getClockThumbnail(id: ClockId): Drawable? = thumbnailCallbacks[id]!!()
+        override fun getClockThumbnail(id: ClockId): Drawable? = thumbnailCallbacks[id]!!(id)
 
         fun addClock(
             id: ClockId,
             name: String,
-            create: () -> ClockController = ::failFactory,
+            create: (ClockId) -> ClockController = ::failFactory,
-            getThumbnail: () -> Drawable? = ::failThumbnail
+            getThumbnail: (ClockId) -> Drawable? = ::failThumbnail
         ): FakeClockPlugin {
             metadata.add(ClockMetadata(id, name))
             createCallbacks[id] = create
@@ -99,7 +104,8 @@ class ClockRegistryTest : SysuiTestCase() {
 
     @Before
     fun setUp() {
-        dispatcher = StandardTestDispatcher()
+        scheduler = TestCoroutineScheduler()
+        dispatcher = StandardTestDispatcher(scheduler)
         scope = TestScope(dispatcher)
 
         fakeDefaultProvider = FakeClockPlugin()
@@ -116,6 +122,8 @@ class ClockRegistryTest : SysuiTestCase() {
             isEnabled = true,
             handleAllUsers = true,
             defaultClockProvider = fakeDefaultProvider,
+            keepAllLoaded = true,
+            subTag = "Test",
         ) {
             override fun querySettings() { }
             override fun applySettings(value: ClockSettings?) {
@@ -142,8 +150,8 @@ class ClockRegistryTest : SysuiTestCase() {
             .addClock("clock_3", "clock 3")
             .addClock("clock_4", "clock 4")
 
-        pluginListener.onPluginConnected(plugin1, mockContext)
+        pluginListener.onPluginLoaded(plugin1, mockContext, mockPluginLifecycle)
-        pluginListener.onPluginConnected(plugin2, mockContext)
+        pluginListener.onPluginLoaded(plugin2, mockContext, mockPluginLifecycle)
         val list = registry.getClocks()
         assertEquals(
             list,
@@ -165,16 +173,18 @@ class ClockRegistryTest : SysuiTestCase() {
 
     @Test
     fun clockIdConflict_ErrorWithoutCrash() {
+        val mockPluginLifecycle1 = mock<PluginLifecycleManager<ClockProviderPlugin>>()
         val plugin1 = FakeClockPlugin()
             .addClock("clock_1", "clock 1", { mockClock }, { mockThumbnail })
             .addClock("clock_2", "clock 2", { mockClock }, { mockThumbnail })
 
+        val mockPluginLifecycle2 = mock<PluginLifecycleManager<ClockProviderPlugin>>()
         val plugin2 = FakeClockPlugin()
             .addClock("clock_1", "clock 1")
             .addClock("clock_2", "clock 2")
 
-        pluginListener.onPluginConnected(plugin1, mockContext)
+        pluginListener.onPluginLoaded(plugin1, mockContext, mockPluginLifecycle1)
-        pluginListener.onPluginConnected(plugin2, mockContext)
+        pluginListener.onPluginLoaded(plugin2, mockContext, mockPluginLifecycle2)
         val list = registry.getClocks()
         assertEquals(
             list,
@@ -202,8 +212,8 @@ class ClockRegistryTest : SysuiTestCase() {
             .addClock("clock_4", "clock 4")
 
         registry.applySettings(ClockSettings("clock_3", null))
-        pluginListener.onPluginConnected(plugin1, mockContext)
+        pluginListener.onPluginLoaded(plugin1, mockContext, mockPluginLifecycle)
-        pluginListener.onPluginConnected(plugin2, mockContext)
+        pluginListener.onPluginLoaded(plugin2, mockContext, mockPluginLifecycle)
 
         val clock = registry.createCurrentClock()
         assertEquals(mockClock, clock)
@@ -220,9 +230,9 @@ class ClockRegistryTest : SysuiTestCase() {
             .addClock("clock_4", "clock 4")
 
         registry.applySettings(ClockSettings("clock_3", null))
-        pluginListener.onPluginConnected(plugin1, mockContext)
+        pluginListener.onPluginLoaded(plugin1, mockContext, mockPluginLifecycle)
-        pluginListener.onPluginConnected(plugin2, mockContext)
+        pluginListener.onPluginLoaded(plugin2, mockContext, mockPluginLifecycle)
-        pluginListener.onPluginDisconnected(plugin2)
+        pluginListener.onPluginUnloaded(plugin2, mockPluginLifecycle)
 
         val clock = registry.createCurrentClock()
         assertEquals(clock, mockDefaultClock)
@@ -230,15 +240,16 @@ class ClockRegistryTest : SysuiTestCase() {
 
     @Test
     fun pluginRemoved_clockAndListChanged() {
+        val mockPluginLifecycle1 = mock<PluginLifecycleManager<ClockProviderPlugin>>()
         val plugin1 = FakeClockPlugin()
             .addClock("clock_1", "clock 1")
             .addClock("clock_2", "clock 2")
 
+        val mockPluginLifecycle2 = mock<PluginLifecycleManager<ClockProviderPlugin>>()
         val plugin2 = FakeClockPlugin()
             .addClock("clock_3", "clock 3", { mockClock })
             .addClock("clock_4", "clock 4")
 
-
         var changeCallCount = 0
         var listChangeCallCount = 0
         registry.registerClockChangeListener(object : ClockRegistry.ClockChangeListener {
@@ -247,23 +258,38 @@ class ClockRegistryTest : SysuiTestCase() {
         })
 
         registry.applySettings(ClockSettings("clock_3", null))
-        assertEquals(0, changeCallCount)
+        scheduler.runCurrent()
+        assertEquals(1, changeCallCount)
         assertEquals(0, listChangeCallCount)
 
-        pluginListener.onPluginConnected(plugin1, mockContext)
+        pluginListener.onPluginLoaded(plugin1, mockContext, mockPluginLifecycle1)
-        assertEquals(0, changeCallCount)
+        scheduler.runCurrent()
+        assertEquals(1, changeCallCount)
         assertEquals(1, listChangeCallCount)
 
-        pluginListener.onPluginConnected(plugin2, mockContext)
+        pluginListener.onPluginLoaded(plugin2, mockContext, mockPluginLifecycle2)
-        assertEquals(1, changeCallCount)
+        scheduler.runCurrent()
+        assertEquals(2, changeCallCount)
         assertEquals(2, listChangeCallCount)
 
-        pluginListener.onPluginDisconnected(plugin1)
+        pluginListener.onPluginUnloaded(plugin1, mockPluginLifecycle1)
-        assertEquals(1, changeCallCount)
+        scheduler.runCurrent()
+        assertEquals(2, changeCallCount)
+        assertEquals(2, listChangeCallCount)
+
+        pluginListener.onPluginUnloaded(plugin2, mockPluginLifecycle2)
+        scheduler.runCurrent()
+        assertEquals(3, changeCallCount)
+        assertEquals(2, listChangeCallCount)
+
+        pluginListener.onPluginDetached(mockPluginLifecycle1)
+        scheduler.runCurrent()
+        assertEquals(3, changeCallCount)
         assertEquals(3, listChangeCallCount)
 
-        pluginListener.onPluginDisconnected(plugin2)
+        pluginListener.onPluginDetached(mockPluginLifecycle2)
-        assertEquals(2, changeCallCount)
+        scheduler.runCurrent()
+        assertEquals(3, changeCallCount)
         assertEquals(4, listChangeCallCount)
     }
 

packages/SystemUI/tests/src/com/android/systemui/shared/plugins/PluginActionManagerTest.java

@@ -79,11 +79,11 @@ public class PluginActionManagerTest extends SysuiTestCase {
     private PluginInstance<TestPlugin> mPluginInstance;
     private PluginInstance.Factory mPluginInstanceFactory = new PluginInstance.Factory(
             this.getClass().getClassLoader(),
-            new PluginInstance.InstanceFactory<>(), new PluginInstance.VersionChecker(),
+            new PluginInstance.InstanceFactory<>(), new PluginInstance.VersionCheckerImpl(),
             Collections.emptyList(), false) {
         @Override
         public <T extends Plugin> PluginInstance<T> create(Context context, ApplicationInfo appInfo,
-                ComponentName componentName, Class<T> pluginClass) {
+                ComponentName componentName, Class<T> pluginClass, PluginListener<T> listener) {
             return (PluginInstance<T>) mPluginInstance;
         }
     };
@@ -128,7 +128,7 @@ public class PluginActionManagerTest extends SysuiTestCase {
         createPlugin();
 
         // Verify startup lifecycle
-        verify(mPluginInstance).onCreate(mContext, mMockListener);
+        verify(mPluginInstance).onCreate();
     }
 
     @Test
@@ -140,7 +140,7 @@ public class PluginActionManagerTest extends SysuiTestCase {
         mFakeExecutor.runAllReady();
 
         // Verify shutdown lifecycle
-        verify(mPluginInstance).onDestroy(mMockListener);
+        verify(mPluginInstance).onDestroy();
     }
 
     @Test
@@ -152,9 +152,9 @@ public class PluginActionManagerTest extends SysuiTestCase {
         mFakeExecutor.runAllReady();
 
         // Verify the old one was destroyed.
-        verify(mPluginInstance).onDestroy(mMockListener);
+        verify(mPluginInstance).onDestroy();
         verify(mPluginInstance, Mockito.times(2))
-                .onCreate(mContext, mMockListener);
+                .onCreate();
     }
 
     @Test
@@ -188,7 +188,7 @@ public class PluginActionManagerTest extends SysuiTestCase {
         mFakeExecutor.runAllReady();
 
         // Verify startup lifecycle
-        verify(mPluginInstance).onCreate(mContext, mMockListener);
+        verify(mPluginInstance).onCreate();
     }
 
     @Test

packages/SystemUI/tests/src/com/android/systemui/shared/plugins/PluginInstanceTest.java

@@ -16,11 +16,9 @@
 
 package com.android.systemui.shared.plugins;
 
+import static junit.framework.Assert.assertEquals;
 import static junit.framework.Assert.assertNotNull;
-
+import static junit.framework.Assert.assertNull;
-import static org.mockito.Matchers.any;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
 
 import android.content.ComponentName;
 import android.content.Context;
@@ -31,6 +29,7 @@ import androidx.test.runner.AndroidJUnit4;
 
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.plugins.Plugin;
+import com.android.systemui.plugins.PluginLifecycleManager;
 import com.android.systemui.plugins.PluginListener;
 import com.android.systemui.plugins.annotations.ProvidesInterface;
 import com.android.systemui.plugins.annotations.Requires;
@@ -38,46 +37,64 @@ import com.android.systemui.plugins.annotations.Requires;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.mockito.Mock;
-import org.mockito.MockitoAnnotations;
 
+import java.lang.ref.WeakReference;
 import java.util.Collections;
+import java.util.concurrent.atomic.AtomicInteger;
 
 @SmallTest
 @RunWith(AndroidJUnit4.class)
 public class PluginInstanceTest extends SysuiTestCase {
 
     private static final String PRIVILEGED_PACKAGE = "com.android.systemui.plugins";
-
+    private static final ComponentName TEST_PLUGIN_COMPONENT_NAME =
-    @Mock
-    private TestPluginImpl mMockPlugin;
-    @Mock
-    private PluginListener<TestPlugin> mMockListener;
-    @Mock
-    private VersionInfo mVersionInfo;
-    ComponentName mTestPluginComponentName =
             new ComponentName(PRIVILEGED_PACKAGE, TestPluginImpl.class.getName());
+
+    private FakeListener mPluginListener;
+    private VersionInfo mVersionInfo;
+    private VersionInfo.InvalidVersionException mVersionException;
+    private PluginInstance.VersionChecker mVersionChecker;
+
+    private RefCounter mCounter;
     private PluginInstance<TestPlugin> mPluginInstance;
     private PluginInstance.Factory mPluginInstanceFactory;
-
     private ApplicationInfo mAppInfo;
-    private Context mPluginContext;
+
-    @Mock
+    // Because we're testing memory in this file, we must be careful not to assert the target
-    private PluginInstance.VersionChecker mVersionChecker;
+    // objects, or capture them via mockito if we expect the garbage collector to later free them.
+    // Both JUnit and Mockito will save references and prevent these objects from being cleaned up.
+    private WeakReference<TestPluginImpl> mPlugin;
+    private WeakReference<Context> mPluginContext;
 
     @Before
     public void setup() throws Exception {
-        MockitoAnnotations.initMocks(this);
+        mCounter = new RefCounter();
         mAppInfo = mContext.getApplicationInfo();
-        mAppInfo.packageName = mTestPluginComponentName.getPackageName();
+        mAppInfo.packageName = TEST_PLUGIN_COMPONENT_NAME.getPackageName();
-        when(mVersionChecker.checkVersion(any(), any(), any())).thenReturn(mVersionInfo);
+        mPluginListener = new FakeListener();
+        mVersionInfo = new VersionInfo();
+        mVersionChecker = new PluginInstance.VersionChecker() {
+            @Override
+            public <T extends Plugin> VersionInfo checkVersion(
+                    Class<T> instanceClass,
+                    Class<T> pluginClass,
+                    Plugin plugin
+            ) {
+                if (mVersionException != null) {
+                    throw mVersionException;
+                }
+                return mVersionInfo;
+            }
+        };
 
         mPluginInstanceFactory = new PluginInstance.Factory(
                 this.getClass().getClassLoader(),
                 new PluginInstance.InstanceFactory<TestPlugin>() {
                     @Override
                     TestPlugin create(Class cls) {
-                        return mMockPlugin;
+                        TestPluginImpl plugin = new TestPluginImpl(mCounter);
+                        mPlugin = new WeakReference<>(plugin);
+                        return plugin;
                     }
                 },
                 mVersionChecker,
@@ -85,8 +102,9 @@ public class PluginInstanceTest extends SysuiTestCase {
                 false);
 
         mPluginInstance = mPluginInstanceFactory.create(
-                mContext, mAppInfo, mTestPluginComponentName, TestPlugin.class);
+                mContext, mAppInfo, TEST_PLUGIN_COMPONENT_NAME,
-        mPluginContext = mPluginInstance.getPluginContext();
+                TestPlugin.class, mPluginListener);
+        mPluginContext = new WeakReference<>(mPluginInstance.getPluginContext());
     }
 
     @Test
@@ -96,29 +114,51 @@ public class PluginInstanceTest extends SysuiTestCase {
 
     @Test(expected = VersionInfo.InvalidVersionException.class)
     public void testIncorrectVersion() throws Exception {
-
         ComponentName wrongVersionTestPluginComponentName =
                 new ComponentName(PRIVILEGED_PACKAGE, TestPlugin.class.getName());
 
-        when(mVersionChecker.checkVersion(any(), any(), any())).thenThrow(
+        mVersionException = new VersionInfo.InvalidVersionException("test", true);
-                new VersionInfo.InvalidVersionException("test", true));
 
         mPluginInstanceFactory.create(
-                mContext, mAppInfo, wrongVersionTestPluginComponentName, TestPlugin.class);
+                mContext, mAppInfo, wrongVersionTestPluginComponentName,
+                TestPlugin.class, mPluginListener);
     }
 
     @Test
     public void testOnCreate() {
-        mPluginInstance.onCreate(mContext, mMockListener);
+        mPluginInstance.onCreate();
-        verify(mMockPlugin).onCreate(mContext, mPluginContext);
+        assertEquals(1, mPluginListener.mAttachedCount);
-        verify(mMockListener).onPluginConnected(mMockPlugin, mPluginContext);
+        assertEquals(1, mPluginListener.mLoadCount);
+        assertEquals(mPlugin.get(), mPluginInstance.getPlugin());
+        assertInstances(1, 1);
     }
 
     @Test
     public void testOnDestroy() {
-        mPluginInstance.onDestroy(mMockListener);
+        mPluginInstance.onDestroy();
-        verify(mMockListener).onPluginDisconnected(mMockPlugin);
+        assertEquals(1, mPluginListener.mDetachedCount);
-        verify(mMockPlugin).onDestroy();
+        assertEquals(1, mPluginListener.mUnloadCount);
+        assertNull(mPluginInstance.getPlugin());
+        assertInstances(0, -1); // Destroyed but never created
+    }
+
+    @Test
+    public void testOnRepeatedlyLoadUnload_PluginFreed() {
+        mPluginInstance.onCreate();
+        mPluginInstance.loadPlugin();
+        assertInstances(1, 1);
+
+        mPluginInstance.unloadPlugin();
+        assertNull(mPluginInstance.getPlugin());
+        assertInstances(0, 0);
+
+        mPluginInstance.loadPlugin();
+        assertInstances(1, 1);
+
+        mPluginInstance.unloadPlugin();
+        mPluginInstance.onDestroy();
+        assertNull(mPluginInstance.getPlugin());
+        assertInstances(0, 0);
     }
 
     // This target class doesn't matter, it just needs to have a Requires to hit the flow where
@@ -129,10 +169,103 @@ public class PluginInstanceTest extends SysuiTestCase {
         String ACTION = "testAction";
     }
 
+    public void assertInstances(Integer allocated, Integer created) {
+        // Run the garbage collector to finalize and deallocate outstanding
+        // instances. Since the GC doesn't always appear to want to run
+        // completely when we ask, we ask it 10 times in a short loop.
+        for (int i = 0; i < 10; i++) {
+            System.runFinalization();
+            System.gc();
+        }
+
+        mCounter.assertInstances(allocated, created);
+    }
+
+    public static class RefCounter {
+        public final AtomicInteger mAllocatedInstances = new AtomicInteger();
+        public final AtomicInteger mCreatedInstances = new AtomicInteger();
+
+        public void assertInstances(Integer allocated, Integer created) {
+            if (allocated != null) {
+                assertEquals(allocated.intValue(), mAllocatedInstances.get());
+            }
+            if (created != null) {
+                assertEquals(created.intValue(), mCreatedInstances.get());
+            }
+        }
+    }
+
     @Requires(target = TestPlugin.class, version = TestPlugin.VERSION)
     public static class TestPluginImpl implements TestPlugin {
+        public final RefCounter mCounter;
+        public TestPluginImpl(RefCounter counter) {
+            mCounter = counter;
+            mCounter.mAllocatedInstances.getAndIncrement();
+        }
+
+        @Override
+        public void finalize() {
+            mCounter.mAllocatedInstances.getAndDecrement();
+        }
+
         @Override
         public void onCreate(Context sysuiContext, Context pluginContext) {
+            mCounter.mCreatedInstances.getAndIncrement();
+        }
+
+        @Override
+        public void onDestroy() {
+            mCounter.mCreatedInstances.getAndDecrement();
+        }
+    }
+
+    public class FakeListener implements PluginListener<TestPlugin> {
+        public int mAttachedCount = 0;
+        public int mDetachedCount = 0;
+        public int mLoadCount = 0;
+        public int mUnloadCount = 0;
+
+        @Override
+        public void onPluginAttached(PluginLifecycleManager<TestPlugin> manager) {
+            mAttachedCount++;
+            assertEquals(PluginInstanceTest.this.mPluginInstance, manager);
+        }
+
+        @Override
+        public void onPluginDetached(PluginLifecycleManager<TestPlugin> manager) {
+            mDetachedCount++;
+            assertEquals(PluginInstanceTest.this.mPluginInstance, manager);
+        }
+
+        @Override
+        public void onPluginLoaded(
+                TestPlugin plugin,
+                Context pluginContext,
+                PluginLifecycleManager<TestPlugin> manager
+        ) {
+            mLoadCount++;
+            TestPlugin expectedPlugin = PluginInstanceTest.this.mPlugin.get();
+            if (expectedPlugin != null) {
+                assertEquals(expectedPlugin, plugin);
+            }
+            Context expectedContext = PluginInstanceTest.this.mPluginContext.get();
+            if (expectedContext != null) {
+                assertEquals(expectedContext, pluginContext);
+            }
+            assertEquals(PluginInstanceTest.this.mPluginInstance, manager);
+        }
+
+        @Override
+        public void onPluginUnloaded(
+                TestPlugin plugin,
+                PluginLifecycleManager<TestPlugin> manager
+        ) {
+            mUnloadCount++;
+            TestPlugin expectedPlugin = PluginInstanceTest.this.mPlugin.get();
+            if (expectedPlugin != null) {
+                assertEquals(expectedPlugin, plugin);
+            }
+            assertEquals(PluginInstanceTest.this.mPluginInstance, manager);
         }
     }
 }