feat/1206 - SDK: Adding events to shielded sync

apps/namadillo/src/App/Masp/ShieldedBalanceChart.tsx

@@ -1,13 +1,47 @@
 import { Heading, PieChart, SkeletonLoading } from "@namada/components";
+import { ProgressBarNames, SdkEvents } from "@namada/shared";
 import { AtomErrorBoundary } from "App/Common/AtomErrorBoundary";
 import { FiatCurrency } from "App/Common/FiatCurrency";
-import { shieldedTokensAtom } from "atoms/balance/atoms";
+import { shieldedSyncAtom, shieldedTokensAtom } from "atoms/balance/atoms";
 import { getTotalDollar } from "atoms/balance/functions";
 import { useAtomValue } from "jotai";
+import { useEffect, useState } from "react";
 import { colors } from "theme";
 
 export const ShieldedBalanceChart = (): JSX.Element => {
   const shieldedTokensQuery = useAtomValue(shieldedTokensAtom);
+  const shieldedSyncProgress = useAtomValue(shieldedSyncAtom);
+
+  const [progress, setProgress] = useState({
+    [ProgressBarNames.Scanned]: 0,
+    [ProgressBarNames.Fetched]: 0,
+    [ProgressBarNames.Applied]: 0,
+  });
+
+  useEffect(() => {
+    shieldedSyncProgress &&
+      shieldedSyncProgress.on(
+        SdkEvents.ProgressBarIncremented,
+        ({ name, current, total }) => {
+          const perc =
+            total === 0 ? 0 : (
+              Math.min(Math.floor((current / total) * 100), 100)
+            );
+          setProgress((prev) => ({
+            ...prev,
+            [name]: perc,
+          }));
+        }
+      );
+
+    shieldedSyncProgress &&
+      shieldedSyncProgress.on(SdkEvents.ProgressBarFinished, ({ name }) => {
+        setProgress((prev) => ({
+          ...prev,
+          [name]: 100,
+        }));
+      });
+  }, [shieldedSyncProgress]);
 
   const shieldedDollars = getTotalDollar(shieldedTokensQuery.data);
 
@@ -48,10 +82,16 @@ export const ShieldedBalanceChart = (): JSX.Element => {
                 </div>
               </PieChart>
               <div className="absolute -bottom-4 -left-2 text-[10px]">
-                *Balances exclude NAM until phase 5
+                *Balances exclude NAM until phase 5{" "}
               </div>
             </>
           }
+          <div className="absolute top-0 right-0 text-right">
+            Shieled sync progress: <br />
+            Scanned: {progress[ProgressBarNames.Scanned]}% <br />
+            Fetched: {progress[ProgressBarNames.Fetched]}% <br />
+            Applied: {progress[ProgressBarNames.Applied]}%
+          </div>
         </AtomErrorBoundary>
       </div>
     </div>

apps/namadillo/src/atoms/balance/atoms.ts

@@ -1,4 +1,5 @@
 import { Asset } from "@chain-registry/types";
+import { SdkEvents } from "@namada/shared";
 import {
   accountsAtom,
   defaultAccountAtom,
@@ -11,8 +12,11 @@ import {
 } from "atoms/chain";
 import { shouldUpdateBalanceAtom } from "atoms/etc";
 import { availableAssetsAtom } from "atoms/integrations";
+import { maspIndexerUrlAtom, rpcUrlAtom } from "atoms/settings";
 import { queryDependentFn } from "atoms/utils";
 import BigNumber from "bignumber.js";
+import EventEmitter from "events";
+import { atom } from "jotai";
 import { atomWithQuery } from "jotai-tanstack-query";
 import { namadaAsset } from "registry/namadaAsset";
 import { AddressWithAsset } from "types";
@@ -21,50 +25,94 @@ import {
   mapNamadaAddressesToAssets,
   mapNamadaAssetsToTokenBalances,
 } from "./functions";
-import { fetchCoinPrices, fetchShieldedBalance } from "./services";
+import {
+  fetchCoinPrices,
+  fetchShieldedBalance,
+  shieldedSync,
+} from "./services";
 
 export type TokenBalance = AddressWithAsset & {
   balance: BigNumber;
   dollar?: BigNumber;
 };
 
-export const viewingKeyAtom = atomWithQuery<string>((get) => {
+export const viewingKeysAtom = atomWithQuery<[string, string[]]>((get) => {
   const accountsQuery = get(accountsAtom);
   const defaultAccountQuery = get(defaultAccountAtom);
 
   return {
-    queryKey: ["viewing-key", accountsQuery.data, defaultAccountQuery.data],
+    queryKey: ["viewing-keys", accountsQuery.data, defaultAccountQuery.data],
     ...queryDependentFn(async () => {
-      const shieldedAccount = accountsQuery.data?.find(
-        (a) => a.isShielded && a.alias === defaultAccountQuery.data?.alias
+      const shieldedAccounts = accountsQuery.data?.filter((a) => a.isShielded);
+      const defaultShieldedAccount = shieldedAccounts?.find(
+        (a) => a.alias === defaultAccountQuery.data?.alias
       );
-      return shieldedAccount?.viewingKey ?? "";
+      const defaultViewingKey = defaultShieldedAccount?.viewingKey ?? "";
+      const viewingKeys =
+        shieldedAccounts?.map((a) => a.viewingKey ?? "") ?? [];
+
+      return [defaultViewingKey, viewingKeys];
     }, [accountsQuery, defaultAccountQuery]),
   };
 });
 
+export const shieldedSyncAtom = atom<EventEmitter | undefined>((get) => {
+  const viewingKeysQuery = get(viewingKeysAtom);
+  const namTokenAddressQuery = get(nativeTokenAddressAtom);
+  const rpcUrl = get(rpcUrlAtom);
+  const maspIndexerUrl = get(maspIndexerUrlAtom);
+
+  const viewingKeys = viewingKeysQuery.data;
+  if (!viewingKeys) {
+    return;
+  }
+  const namTokenAddress = namTokenAddressQuery.data;
+  const [_, allViewingKeys] = viewingKeys!;
+
+  const emitter = shieldedSync(
+    rpcUrl,
+    maspIndexerUrl,
+    namTokenAddress!,
+    allViewingKeys
+  );
+
+  return emitter;
+});
+
 export const shieldedBalanceAtom = atomWithQuery<
   { address: string; amount: BigNumber }[]
 >((get) => {
   const enablePolling = get(shouldUpdateBalanceAtom);
-  const viewingKeyQuery = get(viewingKeyAtom);
+  const viewingKeysQuery = get(viewingKeysAtom);
   const tokenAddressesQuery = get(tokenAddressesAtom);
   const namTokenAddressQuery = get(nativeTokenAddressAtom);
+  const rpcUrl = get(rpcUrlAtom);
+  const maspIndexerUrl = get(maspIndexerUrlAtom);
+  const shieldedSync = get(shieldedSyncAtom);
 
   return {
     refetchInterval: enablePolling ? 1000 : false,
     queryKey: [
       "shielded-balance",
-      viewingKeyQuery.data,
+      viewingKeysQuery.data,
       tokenAddressesQuery.data,
       namTokenAddressQuery.data,
+      rpcUrl,
+      maspIndexerUrl,
+      shieldedSync,
     ],
     ...queryDependentFn(async () => {
-      const viewingKey = viewingKeyQuery.data;
+      const viewingKeys = viewingKeysQuery.data;
       const tokenAddresses = tokenAddressesQuery.data;
-      if (!viewingKey || !tokenAddresses) {
+      if (!viewingKeys || !tokenAddresses) {
         return [];
       }
+      const [viewingKey] = viewingKeys;
+
+      await new Promise<void>((resolve) => {
+        shieldedSync!.once(SdkEvents.ProgressBarFinished, () => resolve());
+      });
+
       const response = await fetchShieldedBalance(
         viewingKey,
         tokenAddresses.map((t) => t.address)
@@ -78,7 +126,7 @@ export const shieldedBalanceAtom = atomWithQuery<
           : new BigNumber(amount),
       }));
       return shieldedBalance;
-    }, [viewingKeyQuery, tokenAddressesQuery, namTokenAddressQuery]),
+    }, [viewingKeysQuery, tokenAddressesQuery, namTokenAddressQuery]),
   };
 });
 

apps/namadillo/src/atoms/balance/services.ts

@@ -1,5 +1,13 @@
-import { Balance } from "@namada/sdk/web";
+import * as Comlink from "comlink";
+import { EventEmitter } from "events";
+
+import { Balance, SdkEvents } from "@namada/sdk/web";
 import { getSdkInstance } from "utils/sdk";
+import {
+  Events,
+  Worker as ShieldedSyncWorkerApi,
+} from "workers/ShieldedSyncWorker";
+import ShieldedSyncWorker from "workers/ShieldedSyncWorker?worker";
 
 const sqsOsmosisApi = "https://sqs.osmosis.zone/";
 
@@ -11,6 +19,44 @@ export const fetchCoinPrices = async (
     `${sqsOsmosisApi}/tokens/prices?base=${assetBaseList.sort((a, b) => a.localeCompare(b)).join(",")}`
   ).then((res) => res.json());
 
+type ShieldedSyncEventMap = Record<
+  | SdkEvents.ProgressBarStarted
+  | SdkEvents.ProgressBarIncremented
+  | SdkEvents.ProgressBarFinished,
+  Events[]
+>;
+
+export function shieldedSync(
+  rpcUrl: string,
+  maspIndexerUrl: string,
+  token: string,
+  viewingKeys: string[]
+): EventEmitter<ShieldedSyncEventMap> {
+  const worker = new ShieldedSyncWorker();
+  const shieldedSyncWorker = Comlink.wrap<ShieldedSyncWorkerApi>(worker);
+  const emitter = new EventEmitter<ShieldedSyncEventMap>();
+
+  worker.onmessage = (event: MessageEvent<Events>) => {
+    emitter.emit(event.data.type, event.data);
+  };
+
+  (async () => {
+    await shieldedSyncWorker.init({
+      type: "init",
+      payload: { rpcUrl, maspIndexerUrl, token },
+    });
+
+    await shieldedSyncWorker.sync({
+      type: "sync",
+      payload: { vks: viewingKeys },
+    });
+
+    worker.terminate();
+  })();
+
+  return emitter;
+}
+
 export const fetchShieldedBalance = async (
   viewingKey: string,
   addresses: string[]
@@ -19,7 +65,6 @@ export const fetchShieldedBalance = async (
   // return await mockShieldedBalance(viewingKey);
 
   const sdk = await getSdkInstance();
-  await sdk.rpc.shieldedSync([viewingKey]);
   return await sdk.rpc.queryBalance(viewingKey, addresses);
 };
 

apps/namadillo/src/workers/ShieldedSyncWorker.ts

@@ -1,14 +1,54 @@
 import { initMulticore } from "@namada/sdk/inline-init";
-import { getSdk, Sdk } from "@namada/sdk/web";
+import { getSdk, Sdk, SdkEvents } from "@namada/sdk/web";
 import * as Comlink from "comlink";
 import { Init, InitDone, Sync, SyncDone } from "./ShieldedSyncMessages";
 
+export type ProgressBarStarted = {
+  type: SdkEvents.ProgressBarStarted;
+  name: string;
+};
+export type ProgressBarIncremented = {
+  type: SdkEvents.ProgressBarIncremented;
+  name: string;
+  current: number;
+  total: number;
+};
+
+export type ProgressBarFinished = {
+  type: SdkEvents.ProgressBarFinished;
+  name: string;
+};
+export type Events =
+  | ProgressBarStarted
+  | ProgressBarIncremented
+  | ProgressBarFinished;
+
 export class Worker {
   private sdk: Sdk | undefined;
 
   async init(m: Init): Promise<InitDone> {
     const { cryptoMemory } = await initMulticore();
     this.sdk = newSdk(cryptoMemory, m.payload);
+
+    // TODO: this can be reduced to one event listener
+    addEventListener(SdkEvents.ProgressBarStarted, (e) => {
+      const event = e as CustomEvent<string>;
+      const payload = JSON.parse(event.detail);
+      postMessage({ ...payload, type: SdkEvents.ProgressBarStarted });
+    });
+
+    addEventListener(SdkEvents.ProgressBarIncremented, (e) => {
+      const event = e as CustomEvent<string>;
+      const payload = JSON.parse(event.detail);
+      postMessage({ ...payload, type: SdkEvents.ProgressBarIncremented });
+    });
+
+    addEventListener(SdkEvents.ProgressBarFinished, (e) => {
+      const event = e as CustomEvent<string>;
+      const payload = JSON.parse(event.detail);
+      postMessage({ ...payload, type: SdkEvents.ProgressBarFinished });
+    });
+
     return { type: "init-done", payload: null };
   }
 

apps/namadillo/vite.config.mjs

@@ -25,6 +25,9 @@ export default defineConfig(() => {
         overlay: { initialIsOpen: false },
       }),
     ],
+    worker: {
+      plugins: () => [tsconfigPaths()],
+    },
     optimizeDeps: {
       esbuildOptions: {
         // Node.js global to browser globalThis

packages/sdk/src/index.ts

@@ -33,7 +33,7 @@ export type {
 export { TxType, TxTypeLabel } from "./tx";
 export type { SupportedTx } from "./tx";
 
-export { Sdk } from "./sdk";
+export { Sdk, SdkEvents } from "./sdk";
 
 export { publicKeyToBech32 } from "./keys";
 

packages/sdk/src/sdk.ts

@@ -10,6 +10,8 @@ import { Rpc } from "./rpc";
 import { Signing } from "./signing";
 import { Tx } from "./tx";
 
+export { SdkEvents } from "@namada/shared";
+
 /**
  * API for interacting with Namada SDK
  */
@@ -114,6 +116,7 @@ export class Sdk {
 
   /**
    * Return SDK Package version
+   * @returns SDK version
    */
   getVersion(): string {
     return packageJson.version;

packages/shared/lib/Cargo.toml

@@ -46,6 +46,11 @@ subtle-encoding = "0.5.1"
 version = "0.3.4"
 features = [
   'console',
+  'Document',
+  'Event',
+  'EventTarget',
+  'CustomEvent',
+  'CustomEventInit',
   'Headers',
   'Request',
   'RequestInit',