WIP: Add support for caching prepared WASM

Author: FranklinWaller

package.json

@@ -2,8 +2,8 @@
     "name": "@fluxprotocol/oracle-vm",
     "description": "Oracle VM for Flux Protocol",
     "version": "1.1.0",
-    "main": "dist/src/main.js",
-    "types": "dist/src/main.d.ts",
+    "main": "dist/Process.js",
+    "types": "dist/Process.d.ts",
     "devDependencies": {
         "@types/big.js": "^6.1.2",
         "@types/bn.js": "^5.1.0",

src/Process.ts

@@ -6,21 +6,24 @@ import { extractMessageFromEvent } from './services/WorkerService';
 import { MessageType, RequestMessage } from './models/WorkerMessage';
 import ProcessMethods from './ProcessMethods';
 import { ExecuteResult } from './models/ExecuteResult';
+import ICache from './models/Cache';
+import InMemoryCache from './models/InMemoryCache';
 
 export default class Process extends EventEmitter {
-    public context: Context;
     private processMethods: ProcessMethods;
     private notifierBuffer = new SharedArrayBuffer(16);
 
-    constructor(context: Context) {
+    constructor(
+        public context: Context, 
+        private cache: ICache = new InMemoryCache()
+    ) {
         super();
-
-        this.context = context;
         this.processMethods = new ProcessMethods(this.notifierBuffer);
     }
 
-    spawn() {
+    async spawn() {
         const worker = new Worker(new URL('./Process.worker.ts', import.meta.url) as NodeURL);
+        const binaryId = this.context.args[0];
 
         worker.addListener('message', async (event: MessageEvent) => {
             const data = extractMessageFromEvent(event);
@@ -30,14 +33,20 @@ export default class Process extends EventEmitter {
                 this.emit('exit', data.value);
             } else if (data.type === MessageType.MethodCall) {
                 this.processMethods.executeMethodCallMessage(data);
+            } else if (data.type === MessageType.Cache) {
+                this.cache.set(binaryId, data.value.binary);
             }
         });
 
+        const cachedBinary = await this.cache.get(binaryId);
+
         const spawnMessage: RequestMessage = {
             type: MessageType.Spawn,
             value: {
                 context: {
                     ...this.context,
+                    isFromCache: cachedBinary ? true : false,
+                    binary: cachedBinary ?? this.context.binary,
                     notifierBuffer: this.notifierBuffer,
                 },
             },
@@ -47,14 +56,16 @@ export default class Process extends EventEmitter {
     }
 }
 
-export function execute(context: Context): Promise<ExecuteResult> {
+export function execute(context: Context, cache: ICache = new InMemoryCache()): Promise<ExecuteResult> {
     return new Promise((resolve) => {
-        const process = new Process(context);
+        const process = new Process(context, cache);
 
         process.on('exit', (result: ExecuteResult) => {
             resolve(result);
         });
 
         process.spawn();
     });
-}
+}
+
+export { ICache, InMemoryCache, Context };

src/Process.worker.ts

@@ -1,6 +1,7 @@
-import { VmContext } from "./models/Context";
+import { VmContext, WorkerContext } from "./models/Context";
 import { MessageType } from "./models/WorkerMessage";
 import { executeWasm } from "./services/VirtualMachineService";
+import { prepareWasmBinary } from "./services/WasmService";
 import { extractMessageFromEvent, workerAddEventListener, selfPostMessage } from "./services/WorkerService";
 
 workerAddEventListener('message', async (event: MessageEvent) => {
@@ -10,7 +11,22 @@ workerAddEventListener('message', async (event: MessageEvent) => {
         return;
     }
 
-    const vmContext = new VmContext(data.value.context);
+    const workerContext: WorkerContext = {
+        ...data.value.context,
+    };
+
+    if (!workerContext.isFromCache) {
+        workerContext.binary = prepareWasmBinary(workerContext.binary);
+
+        selfPostMessage({
+            type: MessageType.Cache,
+            value: {
+                binary: workerContext.binary,
+            },
+        });
+    }
+
+    const vmContext = new VmContext(workerContext);
     const executeResult = await executeWasm(vmContext);
 
     selfPostMessage({

src/models/Cache.ts

@@ -0,0 +1,6 @@
+export default interface ICache {
+    set(id: string, binary: Uint8Array): Promise<void>;
+    get(id: string): Promise<Uint8Array | undefined>;
+    delete(id: string): Promise<void>;
+    clear(): Promise<void>;
+}

src/models/Context.ts

@@ -12,6 +12,7 @@ export interface Context {
 }
 
 export interface WorkerContext extends Context {
+    isFromCache: boolean;
     notifierBuffer: SharedArrayBuffer;
 }
 

src/models/InMemoryCache.ts

@@ -0,0 +1,21 @@
+import ICache from './Cache';
+
+export default class InMemoryCache implements ICache {
+    private internalCache: Map<string, Uint8Array> = new Map();
+
+    async set(id: string, binary: Uint8Array) {
+        this.internalCache.set(id, binary);
+    }
+
+    async delete(id: string) {
+        this.internalCache.delete(id);
+    }
+
+    async get(id: string) {
+        return this.internalCache.get(id);
+    }
+
+    async clear() {
+        this.internalCache.clear();
+    }
+}

src/models/WorkerMessage.ts

@@ -4,6 +4,7 @@ import { ExecuteResult } from "./ExecuteResult";
 export enum MessageType {
     Event = 'EVENT',
     Spawn = 'SPAWN',
+    Cache = 'CACHE',
     Provider = 'PROVIDER',
     Message = 'MESSAGE',
     MethodCall = 'METHOD_CALL',
@@ -35,6 +36,13 @@ export interface MethodCallMessage extends RequestMessageBase {
     };
 }
 
-export type RequestMessage = SpawnMessage | ExitMessage | MethodCallMessage;
+export interface CacheMessage extends RequestMessageBase {
+    type: MessageType.Cache,
+    value: {
+        binary: Uint8Array;
+    }
+}
+
+export type RequestMessage = SpawnMessage | ExitMessage | MethodCallMessage | CacheMessage;
 
 

src/services/VirtualMachineService.ts

@@ -8,16 +8,6 @@ import { generateFixedBuffer } from './CryptoService';
 import VirtualFs from './VirtualFsService';
 import WasmImports from './WasmImports';
 
-const metering = require('wasm-metering');
-
-export function prepareWasmBinary(binary: Uint8Array): Uint8Array {
-    const meteredWasm = metering.meterWASM(binary, {
-        meterType: 'i32',
-    });
-
-    return meteredWasm;
-}
-
 export function executeWasm(context: VmContext): Promise<ExecuteResult> {
     return new Promise(async (resolve) => {
         const wasmFs = new WasmFs();
@@ -47,8 +37,7 @@ export function executeWasm(context: VmContext): Promise<ExecuteResult> {
                 },
             });
 
-            const preparedBinary = prepareWasmBinary(context.binary);
-            const module = await WebAssembly.compile(preparedBinary.buffer);
+            const module = await WebAssembly.compile(context.binary);
             const wasiImports = wasi.getImports(module);
             const wasmImports = new WasmImports(context, wasiImports);
             const instance = await WebAssembly.instantiate(module, wasmImports.getImports());

src/services/WasmService.ts

@@ -0,0 +1,9 @@
+const metering = require('wasm-metering');
+
+export function prepareWasmBinary(binary: Uint8Array): Uint8Array {
+    const meteredWasm = metering.meterWASM(binary, {
+        meterType: 'i32',
+    });
+
+    return meteredWasm;
+}

test/vm.test.ts

@@ -1,9 +1,11 @@
-import { execute } from '../dist/Process';
+import { Context, execute, InMemoryCache } from '../dist/Process';
 import { readFileSync } from 'fs';
 
 jest.setTimeout(10_000);
 
 describe('Process', () => {
+    const memoryCache = new InMemoryCache();
+
     it('should be able to run a wasm file', async () => {
         const wasm = readFileSync(__dirname + '/wasm/cowsay.wasm');
         const result = await execute({
@@ -13,22 +15,23 @@ describe('Process', () => {
             gasLimit: (300_000_000_000_000).toString(),
             randomSeed: '0x012',
             timestamp: 1,
-        });
+        }, memoryCache);
 
         expect(result.gasUsed).toBe('26844206');
         expect(result.code).toBe(0);
     });
 
     it('should throw a gas error when the gas limit has been exceeded', async () => {
         const wasm = readFileSync(__dirname + '/wasm/cowsay.wasm');
+
         const result = await execute({
             args: ['cowsay', 'blah'],
             binary: new Uint8Array(wasm),
             env: {},
             gasLimit: (21_000_000).toString(),
             randomSeed: '0x012',
             timestamp: 1,
-        });
+        }, memoryCache);
 
         expect(result.code).toBe(1);
         expect(result.logs[0]).toBe('ERR_OUT_OF_GAS');
@@ -53,7 +56,7 @@ describe('Process', () => {
             gasLimit: (300_000_000_000_000).toString(),
             randomSeed: '0x012',
             timestamp: new Date().getTime(),
-        });
+        }, memoryCache);
 
         const outcome = JSON.parse(result.logs[result.logs.length - 1]);
         expect(outcome.value).toBe('imposter');
@@ -82,11 +85,11 @@ describe('Process', () => {
             gasLimit: (300_000_000_000_000).toString(),
             randomSeed: '0x012',
             timestamp: new Date().getTime(),
-        });
+        }, memoryCache);
 
         const outcome = JSON.parse(result.logs[result.logs.length - 1]);
 
-        expect(result.gasUsed).toBe('627358915');
+        expect(result.gasUsed).toBe('633190654');
         expect(outcome.value).toBe('70500');
     });
 
@@ -113,11 +116,11 @@ describe('Process', () => {
             gasLimit: (300_000_000_000_000).toString(),
             randomSeed: '0x012',
             timestamp: new Date().getTime(),
-        });
+        }, memoryCache);
 
         const outcome = JSON.parse(result.logs[result.logs.length - 1]);
 
-        expect(result.gasUsed).toBe('627070273');
+        expect(result.gasUsed).toBe('632902012');
         expect(outcome.value).toBe('705000000000000000000000000');
     });
 
@@ -140,10 +143,42 @@ describe('Process', () => {
             gasLimit: (300_000_000_000_000).toString(),
             randomSeed: '0x012',
             timestamp: new Date().getTime(),
-        });
+        }, memoryCache);
 
         const outcome = JSON.parse(result.logs[result.logs.length - 1]);
         expect(result.gasUsed).toBe('23223018');
         expect(outcome.value).toBe('493625367592069900000000000000');
     });
+
+    it('should always use cache when available', async () => {
+        const internalMemoryCache = new InMemoryCache();
+        const wasm = readFileSync(__dirname + '/wasm/basic-fetch.wasm');
+        const setSpy = jest.spyOn(internalMemoryCache, 'set');
+        const getSpy = jest.spyOn(internalMemoryCache, 'get');
+
+        const context: Context = {
+            args: [
+                '0xdeadbeef',
+                JSON.stringify([
+                    {
+                        end_point: 'https://api.coinpaprika.com/v1/coins/btc-bitcoin/ohlcv/historical?start=1630612481&end=1630612781',
+                        source_path: '$[0].close',
+                    },
+                ]),
+                'number',
+                (10e24).toString(),
+            ],
+            binary: new Uint8Array(wasm),
+            env: {},
+            gasLimit: (300_000_000_000_000).toString(),
+            randomSeed: '0x012',
+            timestamp: new Date().getTime()
+        };
+
+        await execute(context, internalMemoryCache);
+        await execute(context, internalMemoryCache);
+
+        expect(getSpy).toHaveBeenCalledTimes(2);
+        expect(setSpy).toHaveBeenCalledTimes(1);
+    });
 });