feat: Implement refcounting and vat termination cleanup (#478)

Author: sirtimid

packages/extension/src/kernel-integration/handlers/collect-garbage.ts

@@ -0,0 +1,30 @@
+import { literal } from '@metamask/superstruct';
+import type { Kernel } from '@ocap/kernel';
+import type { MethodSpec, Handler } from '@ocap/rpc-methods';
+import { EmptyJsonArray } from '@ocap/utils';
+
+export const collectGarbageSpec: MethodSpec<
+  'collectGarbage',
+  EmptyJsonArray,
+  null
+> = {
+  method: 'collectGarbage',
+  params: EmptyJsonArray,
+  result: literal(null),
+};
+
+export type CollectGarbageHooks = { kernel: Pick<Kernel, 'collectGarbage'> };
+
+export const collectGarbageHandler: Handler<
+  'collectGarbage',
+  EmptyJsonArray,
+  null,
+  CollectGarbageHooks
+> = {
+  ...collectGarbageSpec,
+  hooks: { kernel: true },
+  implementation: async ({ kernel }): Promise<null> => {
+    kernel.collectGarbage();
+    return null;
+  },
+};

packages/extension/src/kernel-integration/handlers/index.ts

@@ -1,4 +1,8 @@
 import { clearStateHandler, clearStateSpec } from './clear-state.ts';
+import {
+  collectGarbageHandler,
+  collectGarbageSpec,
+} from './collect-garbage.ts';
 import {
   executeDBQueryHandler,
   executeDBQuerySpec,
@@ -33,6 +37,7 @@ export const handlers = {
   restartVat: restartVatHandler,
   sendVatCommand: sendVatCommandHandler,
   terminateAllVats: terminateAllVatsHandler,
+  collectGarbage: collectGarbageHandler,
   terminateVat: terminateVatHandler,
   updateClusterConfig: updateClusterConfigHandler,
 } as const;
@@ -49,6 +54,7 @@ export const methodSpecs = {
   restartVat: restartVatSpec,
   sendVatCommand: sendVatCommandSpec,
   terminateAllVats: terminateAllVatsSpec,
+  collectGarbage: collectGarbageSpec,
   terminateVat: terminateVatSpec,
   updateClusterConfig: updateClusterConfigSpec,
 } as const;

packages/extension/src/ui/App.module.css

@@ -151,10 +151,18 @@ h6 {
   border: none;
 }
 
+.buttonGray {
+  composes: button;
+  background-color: var(--color-gray-300);
+  color: var(--color-black);
+  border: none;
+}
+
 .button:hover:not(:disabled) {
   background-color: var(--color-gray-800);
   color: var(--color-white);
 }
+
 .textButton {
   padding: 0;
   border: 0;

packages/extension/src/ui/components/KernelControls.tsx

@@ -6,7 +6,8 @@ import { useVats } from '../hooks/useVats.ts';
  * @returns A panel for controlling the kernel.
  */
 export const KernelControls: React.FC = () => {
-  const { terminateAllVats, clearState, reload } = useKernelActions();
+  const { terminateAllVats, collectGarbage, clearState, reload } =
+    useKernelActions();
   const { vats } = useVats();
 
   return (
@@ -16,6 +17,9 @@ export const KernelControls: React.FC = () => {
           Terminate All Vats
         </button>
       )}
+      <button onClick={collectGarbage} className={styles.buttonGray}>
+        Collect Garbage
+      </button>
       <button className={styles.buttonDanger} onClick={clearState}>
         Clear All State
       </button>

packages/extension/src/ui/hooks/useKernelActions.ts

@@ -11,6 +11,7 @@ import { usePanelContext } from '../context/PanelContext.tsx';
 export function useKernelActions(): {
   sendKernelCommand: () => void;
   terminateAllVats: () => void;
+  collectGarbage: () => void;
   clearState: () => void;
   reload: () => void;
   launchVat: (bundleUrl: string, vatName: string) => void;
@@ -42,6 +43,18 @@ export function useKernelActions(): {
       .catch(() => logMessage('Failed to terminate all vats', 'error'));
   }, [callKernelMethod, logMessage]);
 
+  /**
+   * Collects garbage.
+   */
+  const collectGarbage = useCallback(() => {
+    callKernelMethod({
+      method: 'collectGarbage',
+      params: [],
+    })
+      .then(() => logMessage('Garbage collected', 'success'))
+      .catch(() => logMessage('Failed to collect garbage', 'error'));
+  }, [callKernelMethod, logMessage]);
+
   /**
    * Clears the kernel state.
    */
@@ -104,6 +117,7 @@ export function useKernelActions(): {
   return {
     sendKernelCommand,
     terminateAllVats,
+    collectGarbage,
     clearState,
     reload,
     launchVat,

packages/extension/test/e2e/vat-manager.test.ts

@@ -263,4 +263,85 @@ test.describe('Vat Manager', () => {
       minimalClusterConfig.vats.main.parameters.name,
     );
   });
+
+  test('should collect garbage', async () => {
+    await popupPage.click('button:text("Database Inspector")');
+    await expect(messageOutput).toContainText(
+      '{"key":"vats.terminated","value":"[]"}',
+    );
+    const v3Values = [
+      '{"key":"e.nextPromiseId.v3","value":"2"}',
+      '{"key":"e.nextObjectId.v3","value":"1"}',
+      '{"key":"ko3.owner","value":"v3"}',
+      '{"key":"v3.c.ko3","value":"R o+0"}',
+      '{"key":"v3.c.o+0","value":"ko3"}',
+      '{"key":"v3.c.kp3","value":"R p-1"}',
+      '{"key":"v3.c.p-1","value":"kp3"}',
+    ];
+    const v1ko3Values = [
+      '{"key":"v1.c.ko3","value":"R o-2"}',
+      '{"key":"v1.c.o-2","value":"ko3"}',
+      '{"key":"ko3.refCount","value":"2,2"}',
+      '{"key":"kp3.state","value":"fulfilled"}',
+      '{"key":"kp3.value","value"',
+    ];
+    await expect(messageOutput).toContainText(
+      '{"key":"kp3.refCount","value":"2"}',
+    );
+    await expect(messageOutput).toContainText('{"key":"vatConfig.v3","value"');
+    for (const value of v3Values) {
+      await expect(messageOutput).toContainText(value);
+    }
+    for (const value of v1ko3Values) {
+      await expect(messageOutput).toContainText(value);
+    }
+    await popupPage.click('button:text("Vat Manager")');
+    await popupPage.locator('td button:text("Terminate")').last().click();
+    await expect(messageOutput).toContainText('Terminated vat "v3"');
+    await popupPage.locator('[data-testid="clear-logs-button"]').click();
+    await expect(messageOutput).toContainText('');
+    await popupPage.click('button:text("Database Inspector")');
+    await expect(messageOutput).toContainText(
+      '{"key":"vats.terminated","value":"[\\"v3\\"]"}',
+    );
+    await expect(messageOutput).not.toContainText(
+      '{"key":"vatConfig.v3","value"',
+    );
+    for (const value of v3Values) {
+      await expect(messageOutput).toContainText(value);
+    }
+    await popupPage.click('button:text("Vat Manager")');
+    await popupPage.click('button:text("Collect Garbage")');
+    await expect(messageOutput).toContainText('Garbage collected');
+    await popupPage.locator('[data-testid="clear-logs-button"]').click();
+    await expect(messageOutput).toContainText('');
+    await popupPage.click('button:text("Database Inspector")');
+    // v3 is gone
+    for (const value of v3Values) {
+      await expect(messageOutput).not.toContainText(value);
+    }
+    // ko3 reference still exists for v1
+    for (const value of v1ko3Values) {
+      await expect(messageOutput).toContainText(value);
+    }
+    // kp3 reference dropped to 1
+    await expect(messageOutput).toContainText(
+      '{"key":"kp3.refCount","value":"1"}',
+    );
+    await popupPage.click('button:text("Vat Manager")');
+    // delete v1
+    await popupPage.locator('td button:text("Terminate")').first().click();
+    await expect(messageOutput).toContainText('Terminated vat "v1"');
+    await popupPage.click('button:text("Collect Garbage")');
+    await expect(messageOutput).toContainText('Garbage collected');
+    await popupPage.locator('[data-testid="clear-logs-button"]').click();
+    await expect(messageOutput).toContainText('');
+    await popupPage.click('button:text("Database Inspector")');
+    await expect(messageOutput).toContainText(
+      '{"key":"vats.terminated","value":"[]"}',
+    );
+    for (const value of v1ko3Values) {
+      await expect(messageOutput).not.toContainText(value);
+    }
+  });
 });

packages/kernel-test/src/garbage-collection.test.ts

@@ -78,8 +78,8 @@ describe('Garbage Collection', () => {
     const createObjectRef = createObjectData.slots[0] as KRef;
     // Verify initial reference counts from database
     const initialRefCounts = kernelStore.getObjectRefCount(createObjectRef);
-    expect(initialRefCounts.reachable).toBe(1);
-    expect(initialRefCounts.recognizable).toBe(1);
+    expect(initialRefCounts.reachable).toBe(3);
+    expect(initialRefCounts.recognizable).toBe(3);
     // Send the object to the importer vat
     const objectRef = kunser(createObjectData);
     await kernel.queueMessageFromKernel(importerKRef, 'storeImport', [
@@ -95,7 +95,7 @@ describe('Garbage Collection', () => {
     // Check that the object is reachable as a promise from the importer vat
     const importerKref = kernelStore.erefToKref(importerVatId, 'p-1') as KRef;
     expect(kernelStore.hasCListEntry(importerVatId, importerKref)).toBe(true);
-    expect(kernelStore.getRefCount(importerKref)).toBe(1);
+    expect(kernelStore.getRefCount(importerKref)).toBe(2);
     // Use the object
     const useResult = await kernel.queueMessageFromKernel(
       importerKRef,
@@ -118,8 +118,8 @@ describe('Garbage Collection', () => {
 
     // Store initial reference count information
     const initialRefCounts = kernelStore.getObjectRefCount(createObjectRef);
-    expect(initialRefCounts.reachable).toBe(1);
-    expect(initialRefCounts.recognizable).toBe(1);
+    expect(initialRefCounts.reachable).toBe(3);
+    expect(initialRefCounts.recognizable).toBe(3);
 
     // Store the reference in the importer vat
     const objectRef = kunser(createObjectData);
@@ -160,8 +160,8 @@ describe('Garbage Collection', () => {
 
     // Check reference counts after dropImports
     const afterWeakRefCounts = kernelStore.getObjectRefCount(createObjectRef);
-    expect(afterWeakRefCounts.reachable).toBe(0);
-    expect(afterWeakRefCounts.recognizable).toBe(1);
+    expect(afterWeakRefCounts.reachable).toBe(2);
+    expect(afterWeakRefCounts.recognizable).toBe(3);
 
     // Now completely forget the import in the importer vat
     // This should trigger retireImports when GC runs
@@ -171,16 +171,15 @@ describe('Garbage Collection', () => {
     // Schedule another reap
     kernel.reapVats((vatId) => vatId === importerVatId);
 
-    // Run 3 cranks to allow bringOutYourDead to be processed
     for (let i = 0; i < 3; i++) {
       await kernel.queueMessageFromKernel(importerKRef, 'noop', []);
       await waitUntilQuiescent(500);
     }
 
-    // Check reference counts after retireImports (both should be decreased)
+    // Check reference counts after retireImports
     const afterForgetRefCounts = kernelStore.getObjectRefCount(createObjectRef);
-    expect(afterForgetRefCounts.reachable).toBe(0);
-    expect(afterForgetRefCounts.recognizable).toBe(0);
+    expect(afterForgetRefCounts.reachable).toBe(2);
+    expect(afterForgetRefCounts.recognizable).toBe(2);
 
     // Now forget the object in the exporter vat
     // This should trigger retireExports when GC runs

packages/kernel-test/src/utils.ts

@@ -150,3 +150,13 @@ export function parseReplyBody(body: string): unknown {
     return body;
   }
 }
+
+/**
+ * Debug the database.
+ *
+ * @param kernelDatabase - The database to debug.
+ */
+export function logDatabase(kernelDatabase: KernelDatabase): void {
+  const result = kernelDatabase.executeQuery('SELECT * FROM kv');
+  console.log(result);
+}