feat: poc display shieled sync progress

apps/namadillo/src/App/App.tsx

@@ -6,8 +6,6 @@ import { useTransactionCallback } from "hooks/useTransactionCallbacks";
 import { useTransactionNotifications } from "hooks/useTransactionNotifications";
 import { Outlet } from "react-router-dom";
 
-import { SdkEvents } from "@heliaxdev/namada-sdk/web";
-import { useEffect } from "react";
 import { ChainLoader } from "./Setup/ChainLoader";
 
 export const history = createBrowserHistory({ window });
@@ -17,55 +15,6 @@ export function App(): JSX.Element {
   useTransactionNotifications();
   useTransactionCallback();
 
-  // TODO: This useEffect should be removed! Just testing here:
-  useEffect(() => {
-    // eslint-disable-next-line
-    const onProgressBarStarted = (e: any): void => {
-      console.info(SdkEvents.ProgressBarStarted, e);
-    };
-
-    // eslint-disable-next-line
-    const onProgressBarIncremented = (e: any): void => {
-      console.info(SdkEvents.ProgressBarIncremented, e);
-    };
-
-    // eslint-disable-next-line
-    const onProgressBarFinished = (e: any): void => {
-      console.info(SdkEvents.ProgressBarFinished, e);
-    };
-
-    document.onreadystatechange = function() {
-      window.addEventListener(
-        SdkEvents.ProgressBarStarted,
-        onProgressBarStarted
-      );
-
-      window.addEventListener(
-        SdkEvents.ProgressBarIncremented,
-        onProgressBarIncremented
-      );
-      window.addEventListener(
-        SdkEvents.ProgressBarFinished,
-        onProgressBarFinished
-      );
-
-      return () => {
-        window.removeEventListener(
-          SdkEvents.ProgressBarStarted,
-          onProgressBarStarted
-        );
-        window.removeEventListener(
-          SdkEvents.ProgressBarIncremented,
-          onProgressBarIncremented
-        );
-        window.removeEventListener(
-          SdkEvents.ProgressBarFinished,
-          onProgressBarFinished
-        );
-      };
-    };
-  }, []);
-
   return (
     <>
       <Toasts />

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

@@ -1,14 +1,23 @@
 import { Heading, PieChart, SkeletonLoading } from "@namada/components";
 import { AtomErrorBoundary } from "App/Common/AtomErrorBoundary";
 import { FiatCurrency } from "App/Common/FiatCurrency";
-import { shieldedTokensAtom } from "atoms/balance/atoms";
+import {
+  shieldedSyncProgressAtom,
+  shieldedTokensAtom,
+} from "atoms/balance/atoms";
 import { getTotalDollar } from "atoms/balance/functions";
 import { useAtomValue } from "jotai";
 import { colors } from "theme";
 
 export const ShieldedBalanceChart = (): JSX.Element => {
   const shieldedTokensQuery = useAtomValue(shieldedTokensAtom);
-
+  const shieledSyncProgress = useAtomValue(shieldedSyncProgressAtom("scanned"));
+  const shieledSyncProgress2 = useAtomValue(
+    shieldedSyncProgressAtom("fetched")
+  );
+  const shieledSyncProgress3 = useAtomValue(
+    shieldedSyncProgressAtom("applied")
+  );
   const shieldedDollars = getTotalDollar(shieldedTokensQuery.data);
 
   return (
@@ -48,12 +57,44 @@ 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 />
+            <Progress name="Scanned" {...shieledSyncProgress} />
+            <Progress name="Fetched" {...shieledSyncProgress2} />
+            <Progress name="Applied" {...shieledSyncProgress3} />
+          </div>
         </AtomErrorBoundary>
       </div>
     </div>
   );
 };
+
+const Progress = ({
+  name,
+  status,
+  current,
+  total,
+}: {
+  name: string;
+  status: string;
+  current: number;
+  total: number;
+}): JSX.Element => {
+  const perc = current === 0 && total === 0 ? 0 : (current / total) * 100;
+  return (
+    <div>
+      {name}:
+      {status === "pending" ?
+        <div>-</div>
+      : status === "loading" ?
+        <div>{Math.floor(perc)}%</div>
+      : status === "success" ?
+        <div>100%</div>
+      : <div>error</div>}
+    </div>
+  );
+};

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

@@ -11,60 +11,91 @@ 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 { atom } from "jotai";
 import { atomWithQuery } from "jotai-tanstack-query";
+import { atomFamily } from "jotai/utils";
 import { namadaAsset } from "registry/namadaAsset";
 import { AddressWithAsset } from "types";
 import {
   findAssetByToken,
   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 shieldedSyncProgressAtom = atomFamily((_name: string) =>
+  // total is 0 so we do not get NaN
+  atom({ status: "pending", current: 0, total: 0 })
+);
+
 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);
 
   return {
     refetchInterval: enablePolling ? 1000 : false,
     queryKey: [
       "shielded-balance",
-      viewingKeyQuery.data,
+      viewingKeysQuery.data,
       tokenAddressesQuery.data,
       namTokenAddressQuery.data,
+      rpcUrl,
+      maspIndexerUrl,
     ],
     ...queryDependentFn(async () => {
-      const viewingKey = viewingKeyQuery.data;
+      const viewingKeys = viewingKeysQuery.data;
       const tokenAddresses = tokenAddressesQuery.data;
-      if (!viewingKey || !tokenAddresses) {
+      if (!viewingKeys || !tokenAddresses) {
         return [];
       }
+      const namTokenAddress = namTokenAddressQuery.data;
+      const [viewingKey, allViewingKeys] = viewingKeys;
+
+      await shieldedSync(
+        rpcUrl,
+        maspIndexerUrl,
+        namTokenAddress!,
+        allViewingKeys
+      );
+
       const response = await fetchShieldedBalance(
         viewingKey,
         tokenAddresses.map((t) => t.address)
@@ -78,7 +109,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,12 @@
-import { Balance } from "@namada/sdk/web";
+import * as Comlink from "comlink";
+
+import { Balance, SdkEvents } from "@namada/sdk/web";
+import { getDefaultStore } from "jotai";
 import { getSdkInstance } from "utils/sdk";
+import { Worker as ShieldedSyncWorkerApi } from "workers/ShieldedSyncWorker";
+import ShieldedSyncWorker from "workers/ShieldedSyncWorker?worker";
+// TODO: circular dependency
+import { shieldedSyncProgressAtom } from "./atoms";
 
 const sqsOsmosisApi = "https://sqs.osmosis.zone/";
 
@@ -11,6 +18,61 @@ export const fetchCoinPrices = async (
     `${sqsOsmosisApi}/tokens/prices?base=${assetBaseList.sort((a, b) => a.localeCompare(b)).join(",")}`
   ).then((res) => res.json());
 
+export const shieldedSync = async (
+  rpcUrl: string,
+  maspIndexerUrl: string,
+  token: string,
+  viewingKeys: string[]
+): Promise<void> => {
+  const worker = new ShieldedSyncWorker();
+  const shieldedSyncWorker = Comlink.wrap<ShieldedSyncWorkerApi>(worker);
+
+  worker.onmessage = (event) => {
+    const store = getDefaultStore();
+    const { data } = event;
+
+    if (SdkEvents.ProgressBarStarted === data.type) {
+      store.set(shieldedSyncProgressAtom(data.payload.name), {
+        status: "loading",
+        current: 0,
+        total: 0,
+      });
+    } else if (SdkEvents.ProgressBarIncremented === data.type) {
+      const { name, current, total } = data.payload;
+
+      store.set(shieldedSyncProgressAtom(name), {
+        status: "loading",
+        current,
+        total,
+      });
+    } else if (SdkEvents.ProgressBarFinished === data.type) {
+      store.set(shieldedSyncProgressAtom(data.payload.name), {
+        status: "success",
+        current: 0,
+        total: 0,
+      });
+    }
+  };
+
+  await shieldedSyncWorker.init({
+    type: "init",
+    payload: {
+      rpcUrl,
+      maspIndexerUrl,
+      token,
+    },
+  });
+
+  await shieldedSyncWorker.sync({
+    type: "sync",
+    payload: {
+      vks: viewingKeys,
+    },
+  });
+
+  worker.terminate();
+};
+
 export const fetchShieldedBalance = async (
   viewingKey: string,
   addresses: string[]
@@ -19,7 +81,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,41 @@
 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 Events =
+  | { type: SdkEvents.ProgressBarStarted; payload: { name: string } }
+  | {
+      type: SdkEvents.ProgressBarIncremented;
+      payload: { name: string; current: number; total: number };
+    }
+  | { type: SdkEvents.ProgressBarFinished; payload: { name: string } };
+
 export class Worker {
   private sdk: Sdk | undefined;
 
   async init(m: Init): Promise<InitDone> {
     const { cryptoMemory } = await initMulticore();
     this.sdk = newSdk(cryptoMemory, m.payload);
+
+    addEventListener(SdkEvents.ProgressBarStarted, (data) => {
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      const payload = JSON.parse((data as any).detail);
+      postMessage({ type: SdkEvents.ProgressBarStarted, payload });
+    });
+
+    addEventListener(SdkEvents.ProgressBarIncremented, (data) => {
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      const payload = JSON.parse((data as any).detail);
+      postMessage({ type: SdkEvents.ProgressBarIncremented, payload });
+    });
+
+    addEventListener(SdkEvents.ProgressBarFinished, (data) => {
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      const payload = JSON.parse((data as any).detail);
+      postMessage({ type: SdkEvents.ProgressBarFinished, payload });
+    });
+
     return { type: "init-done", payload: null };
   }