feat(runtime-core): add app.onUnmount() for registering cleanup functions (#4619)

close #4516

Author: LinusBorg

Diff for: packages/runtime-core/__tests__/apiCreateApp.spec.ts

@@ -344,6 +344,36 @@ describe('api: createApp', () => {
     ).toHaveBeenWarnedTimes(1)
   })
 
+  test('onUnmount', () => {
+    const cleanup = vi.fn().mockName('plugin cleanup')
+    const PluginA: Plugin = app => {
+      app.provide('foo', 1)
+      app.onUnmount(cleanup)
+    }
+    const PluginB: Plugin = {
+      install: (app, arg1, arg2) => {
+        app.provide('bar', arg1 + arg2)
+        app.onUnmount(cleanup)
+      },
+    }
+
+    const app = createApp({
+      render: () => `Test`,
+    })
+    app.use(PluginA)
+    app.use(PluginB)
+
+    const root = nodeOps.createElement('div')
+    app.mount(root)
+
+    //also can be added after mount
+    app.onUnmount(cleanup)
+
+    app.unmount()
+
+    expect(cleanup).toHaveBeenCalledTimes(3)
+  })
+
   test('config.errorHandler', () => {
     const error = new Error()
     const count = ref(0)

Diff for: packages/runtime-core/src/apiCreateApp.ts

@@ -27,6 +27,7 @@ import { version } from '.'
 import { installAppCompatProperties } from './compat/global'
 import type { NormalizedPropsOptions } from './componentProps'
 import type { ObjectEmitsOptions } from './componentEmits'
+import { ErrorCodes, callWithAsyncErrorHandling } from './errorHandling'
 import type { DefineComponent } from './apiDefineComponent'
 
 export interface App<HostElement = any> {
@@ -50,6 +51,7 @@ export interface App<HostElement = any> {
     namespace?: boolean | ElementNamespace,
   ): ComponentPublicInstance
   unmount(): void
+  onUnmount(cb: () => void): void
   provide<T>(key: InjectionKey<T> | string, value: T): this
 
   /**
@@ -214,6 +216,7 @@ export function createAppAPI<HostElement>(
 
     const context = createAppContext()
     const installedPlugins = new WeakSet()
+    const pluginCleanupFns: Array<() => any> = []
 
     let isMounted = false
 
@@ -366,8 +369,23 @@ export function createAppAPI<HostElement>(
         }
       },
 
+      onUnmount(cleanupFn: () => void) {
+        if (__DEV__ && typeof cleanupFn !== 'function') {
+          warn(
+            `Expected function as first argument to app.onUnmount(), ` +
+              `but got ${typeof cleanupFn}`,
+          )
+        }
+        pluginCleanupFns.push(cleanupFn)
+      },
+
       unmount() {
         if (isMounted) {
+          callWithAsyncErrorHandling(
+            pluginCleanupFns,
+            app._instance,
+            ErrorCodes.APP_UNMOUNT_CLEANUP,
+          )
           render(null, app._container)
           if (__DEV__ || __FEATURE_PROD_DEVTOOLS__) {
             app._instance = null

Diff for: packages/runtime-core/src/errorHandling.ts

@@ -23,6 +23,7 @@ export enum ErrorCodes {
   FUNCTION_REF,
   ASYNC_COMPONENT_LOADER,
   SCHEDULER,
+  APP_UNMOUNT_CLEANUP,
 }
 
 export const ErrorTypeStrings: Record<LifecycleHooks | ErrorCodes, string> = {
@@ -57,6 +58,7 @@ export const ErrorTypeStrings: Record<LifecycleHooks | ErrorCodes, string> = {
   [ErrorCodes.SCHEDULER]:
     'scheduler flush. This is likely a Vue internals bug. ' +
     'Please open an issue at https://github.com/vuejs/core .',
+  [ErrorCodes.APP_UNMOUNT_CLEANUP]: 'app unmount cleanup function',
 }
 
 export type ErrorTypes = LifecycleHooks | ErrorCodes