refactor: shielded sync worker to use comlink

anoma · Oct 8, 2024 · f7697d6 · f7697d6
1 parent 6f596fb
commit f7697d6
Show file tree

Hide file tree

Showing 5 changed files with 102 additions and 86 deletions.
diff --git a/apps/namadillo/src/App/App.tsx b/apps/namadillo/src/App/App.tsx
@@ -18,9 +18,12 @@ import { Outlet } from "react-router-dom";
 import { Shield } from "workers/ShieldMessages";
 import {
   registerTransferHandlers,
-  ShieldWorkerApi,
+  Worker as ShieldWorkerApi,
 } from "workers/ShieldWorker";
 import ShieldWorker from "workers/ShieldWorker?worker";
+import { Worker as ShieldedSyncWorkerApi } from "workers/ShieldedSyncWorker";
+import ShieldedSyncWorker from "workers/ShieldedSyncWorker?worker";
+
 import { ChainLoader } from "./Setup/ChainLoader";
 
 export const history = createBrowserHistory({ window });
@@ -33,24 +36,36 @@ export function App(): JSX.Element {
 
   const rpcUrl = useAtomValue(rpcUrlAtom);
 
-  registerTransferHandlers();
-  const shieldWorker = Comlink.wrap<ShieldWorkerApi>(new ShieldWorker());
   const { data: account } = useAtomValue(defaultAccountAtom);
   const { data: chain } = useAtomValue(chainAtom);
   const { data: token } = useAtomValue(nativeTokenAddressAtom);
   const indexerUrl = useAtomValue(indexerUrlAtom);
-  const shiedlingMsgValue = new ShieldingTransferMsgValue({
-    target:
-      "znam1vue386rsee5c65qsvlm9tgj5lqetz6ejkln3j57utc44w2n8upty57z7lh07myrj3clfxyl9lvn",
-    data: [
-      {
-        source: account?.address || "",
+
+  const shieldedSyncWorker = Comlink.wrap<ShieldedSyncWorkerApi>(
+    new ShieldedSyncWorker()
+  );
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  (window as any).shieldedSync = async () => {
+    await shieldedSyncWorker.init({
+      type: "init",
+      payload: {
+        rpcUrl,
         token: token!,
-        amount: BigNumber(100),
       },
-    ],
-  });
+    });
+
+    await shieldedSyncWorker.sync({
+      type: "sync",
+      payload: {
+        vks: [
+          "zvknam1qvgwgy79qqqqpq88yru6n3f3ugfme002t7272a0ke8zdr2kt80jhnjwmgxkwm7yc6ydp8tfh8lmd28n8hrmcvqszjm3tnytryaf4qhwu645xks4nnx64m3fnpm8yr6hrpd8jtsupyzz4knqleuy7jdjz32jcz9ual56vrf3estg0e6kew0g9aqs4vg2d6n569c78ttqw4zw6mvjkhwfprcc804qt3yewsrxf8l67p87ltnqjtjkr35pfnnxavs9c5wqpr2t2lf3husqn4zvux",
+        ],
+      },
+    });
+  };
 
+  registerTransferHandlers();
+  const shieldWorker = Comlink.wrap<ShieldWorkerApi>(new ShieldWorker());
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   (window as any).buildShieldlingTx = async () => {
     await shieldWorker.init({
@@ -61,6 +76,18 @@ export function App(): JSX.Element {
       },
     });
 
+    const shiedlingMsgValue = new ShieldingTransferMsgValue({
+      target:
+        "znam1vue386rsee5c65qsvlm9tgj5lqetz6ejkln3j57utc44w2n8upty57z7lh07myrj3clfxyl9lvn",
+      data: [
+        {
+          source: account?.address || "",
+          token: token!,
+          amount: BigNumber(100),
+        },
+      ],
+    });
+
     const msg: Shield = {
       type: "shield",
       payload: {

diff --git a/apps/namadillo/src/workers/ShieldMessages.ts b/apps/namadillo/src/workers/ShieldMessages.ts
@@ -10,6 +10,9 @@ type InitPayload = {
   token: string;
 };
 
+export type Init = WebWorkerMessage<"init", InitPayload>;
+export type InitDone = WebWorkerMessage<"init-done", null>;
+
 type ShieldPayload = {
   account: Account;
   gasConfig: {
@@ -20,21 +23,16 @@ type ShieldPayload = {
   chain: ChainSettings;
   indexerUrl: string;
 };
-
-type BroadcastPayload = {
-  encodedTx: EncodedTxData<ShieldingTransferMsgValue>;
-  signedTxs: Uint8Array[];
-};
-
-export type Init = WebWorkerMessage<"init", InitPayload>;
-export type InitDone = WebWorkerMessage<"init-done", null>;
-
 export type Shield = WebWorkerMessage<"shield", ShieldPayload>;
 export type ShieldDone = WebWorkerMessage<
   "shield-done",
   EncodedTxData<ShieldingTransferMsgValue>
 >;
 
+type BroadcastPayload = {
+  encodedTx: EncodedTxData<ShieldingTransferMsgValue>;
+  signedTxs: Uint8Array[];
+};
 export type Broadcast = WebWorkerMessage<"broadcast", BroadcastPayload>;
 export type BroadcastDone = WebWorkerMessage<"broadcast-done", null>;
 

diff --git a/apps/namadillo/src/workers/ShieldWorker.ts b/apps/namadillo/src/workers/ShieldWorker.ts
@@ -9,19 +9,21 @@ import {
   Broadcast,
   BroadcastDone,
   Init,
+  InitDone,
   Shield,
   ShieldDone,
 } from "./ShieldMessages";
 import { registerBNTransferHandler } from "./utils";
 
-// TODO: replace with Class
-const worker = {
-  sdk: undefined as Sdk | undefined,
-  async init(m: Init) {
+export class Worker {
+  private sdk: Sdk | undefined;
+
+  async init(m: Init): Promise<InitDone> {
     const { cryptoMemory } = await initMulticore();
     this.sdk = newSdk(cryptoMemory, m.payload);
     return { type: "init-done", payload: null };
-  },
+  }
+
   async shield(m: Shield): Promise<ShieldDone> {
     if (!this.sdk) {
       throw new Error("SDK is not initialized");
@@ -30,15 +32,16 @@ const worker = {
       type: "shield-done",
       payload: await shield(this.sdk, m.payload),
     };
-  },
+  }
+
   async broadcast(m: Broadcast): Promise<BroadcastDone> {
     if (!this.sdk) {
       throw new Error("SDK is not initialized");
     }
     await broadcast(this.sdk, m.payload);
     return { type: "broadcast-done", payload: null };
-  },
-};
+  }
+}
 
 async function shield(
   sdk: Sdk,
@@ -97,7 +100,5 @@ export const registerTransferHandlers = (): void => {
   registerBNTransferHandler<Broadcast>("broadcast");
 };
 
-export type ShieldWorkerApi = typeof worker;
-
 registerTransferHandlers();
-Comlink.expose(worker);
+Comlink.expose(new Worker());
diff --git a/apps/namadillo/src/workers/ShieldedSyncMessages.ts b/apps/namadillo/src/workers/ShieldedSyncMessages.ts
@@ -0,0 +1,18 @@
+import { WebWorkerMessage } from "./utils";
+
+type InitPayload = {
+  rpcUrl: string;
+  token: string;
+  maspIndexerUrl?: string;
+};
+export type Init = WebWorkerMessage<"init", InitPayload>;
+export type InitDone = WebWorkerMessage<"init-done", null>;
+
+type ShiededSyncPayload = {
+  vks: string[];
+};
+export type Sync = WebWorkerMessage<"sync", ShiededSyncPayload>;
+export type SyncDone = WebWorkerMessage<"sync-done", null>;
+
+export type ShieldMessageIn = Init | Sync;
+export type ShieldMessageOut = InitDone | SyncDone;
diff --git a/apps/namadillo/src/workers/ShieldedSyncWorker.ts b/apps/namadillo/src/workers/ShieldedSyncWorker.ts
@@ -1,66 +1,38 @@
-import init, { initMulticore } from "@heliax/namada-sdk/inline-init";
+import { initMulticore } from "@heliax/namada-sdk/inline-init";
+import * as Comlink from "comlink";
 
-import { getSdk } from "@heliax/namada-sdk/web";
+import { getSdk, Sdk } from "@heliax/namada-sdk/web";
+import { Init, InitDone, Sync, SyncDone } from "./ShieldedSyncMessages";
 
-export type ShiededSyncPayload = {
-  vks: string[];
-  rpcUrl: string;
-  maspIndexerUrl?: string;
-};
+export class Worker {
+  private sdk: Sdk | undefined;
 
-export type ShieldedSyncMulticore = {
-  type: "shielded-sync-multicore";
-  payload: ShiededSyncPayload;
-};
-
-export type ShieldedSyncSinglecore = {
-  type: "shielded-sync-singlecore";
-  payload: ShiededSyncPayload;
-};
-
-export type ShieldedSyncMessageType =
-  | ShieldedSyncMulticore
-  | ShieldedSyncSinglecore;
-
-// TODO: replace with comlink
-self.onmessage = async (e: MessageEvent<ShieldedSyncMessageType>) => {
-  const { type, payload } = e.data;
-
-  switch (type) {
-    case "shielded-sync-singlecore": {
-      const { cryptoMemory } = await init();
-      //eslint-disable-next-line
-      console.log("Syncing with single core");
-      await shieldedSync(cryptoMemory, payload);
-      break;
-    }
+  async init(m: Init): Promise<InitDone> {
+    const { cryptoMemory } = await initMulticore();
+    this.sdk = newSdk(cryptoMemory, m.payload);
+    return { type: "init-done", payload: null };
+  }
 
-    case "shielded-sync-multicore": {
-      const { cryptoMemory } = await initMulticore();
-      //eslint-disable-next-line
-      console.log("Syncing with multicore");
-      await shieldedSync(cryptoMemory, payload);
-      break;
+  async sync(m: Sync): Promise<SyncDone> {
+    if (!this.sdk) {
+      throw new Error("SDK is not initialized");
     }
 
-    default:
-      throw new Error(`Unknown message type: ${type}`);
+    await shieldedSync(this.sdk, m.payload);
+    return { type: "sync-done", payload: null };
   }
-};
+}
 
-async function shieldedSync(
+function newSdk(
   cryptoMemory: WebAssembly.Memory,
-  payload: ShiededSyncPayload
-): Promise<void> {
-  const { rpcUrl, maspIndexerUrl } = payload;
+  payload: Init["payload"]
+): Sdk {
+  const { rpcUrl, token, maspIndexerUrl } = payload;
+  return getSdk(cryptoMemory, rpcUrl, maspIndexerUrl || "", "", token);
+}
 
-  const sdk = getSdk(
-    cryptoMemory,
-    rpcUrl,
-    maspIndexerUrl || "",
-    "",
-    // Not really used, but required by the SDK, as long as it's valid address it's fine
-    "tnam1qxfj3sf6a0meahdu9t6znp05g8zx4dkjtgyn9gfu"
-  );
+async function shieldedSync(sdk: Sdk, payload: Sync["payload"]): Promise<void> {
   await sdk.rpc.shieldedSync(payload.vks);
 }
+
+Comlink.expose(new Worker());