feat/1206 - SDK: Adding events to shielded sync

apps/namadillo/src/App/App.tsx

@@ -7,6 +7,7 @@ import { useTransactionCallback } from "hooks/useTransactionCallbacks";
 import { useTransactionNotifications } from "hooks/useTransactionNotifications";
 import { useTransactionWatcher } from "hooks/useTransactionWatcher";
 import { Outlet } from "react-router-dom";
+
 import { ChainLoader } from "./Setup/ChainLoader";
 
 export const history = createBrowserHistory({ window });

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

@@ -1,13 +1,77 @@
 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 { useAtom, useAtomValue } from "jotai";
+import { useEffect, useState } from "react";
 import { colors } from "theme";
+import {
+  ProgressBarFinished,
+  ProgressBarIncremented,
+} from "workers/ShieldedSyncWorker";
 
 export const ShieldedBalanceChart = (): JSX.Element => {
   const shieldedTokensQuery = useAtomValue(shieldedTokensAtom);
+  const [{ data: shieldedSyncProgress, refetch: shieledSync }] =
+    useAtom(shieldedSyncAtom);
+
+  const [progress, setProgress] = useState({
+    [ProgressBarNames.Scanned]: 0,
+    [ProgressBarNames.Fetched]: 0,
+    [ProgressBarNames.Applied]: 0,
+  });
+
+  useEffect(() => {
+    if (!shieldedSyncProgress) return;
+
+    const onProgressBarIncremented = ({
+      name,
+      current,
+      total,
+    }: ProgressBarIncremented): void => {
+      const perc =
+        total === 0 ? 0 : Math.min(Math.floor((current / total) * 100), 100);
+
+      setProgress((prev) => ({
+        ...prev,
+        [name]: perc,
+      }));
+    };
+
+    const onProgressBarFinished = ({ name }: ProgressBarFinished): void => {
+      setProgress((prev) => ({
+        ...prev,
+        [name]: 100,
+      }));
+    };
+
+    shieldedSyncProgress.on(
+      SdkEvents.ProgressBarIncremented,
+      onProgressBarIncremented
+    );
+
+    shieldedSyncProgress.on(
+      SdkEvents.ProgressBarFinished,
+      onProgressBarFinished
+    );
+
+    return () => {
+      shieldedSyncProgress.off(
+        SdkEvents.ProgressBarIncremented,
+        onProgressBarIncremented
+      );
+      shieldedSyncProgress.off(
+        SdkEvents.ProgressBarFinished,
+        onProgressBarFinished
+      );
+    };
+  }, [shieldedSyncProgress]);
+
+  useEffect(() => {
+    shieledSync();
+  }, []);
 
   const shieldedDollars = getTotalDollar(shieldedTokensQuery.data);
 
@@ -48,10 +112,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,3 +1,4 @@
+import { SdkEvents } from "@namada/shared";
 import { Namada } from "@namada/types";
 import {
   accountsAtom,
@@ -10,6 +11,7 @@ import {
   tokenAddressesAtom,
 } from "atoms/chain";
 import { shouldUpdateBalanceAtom } from "atoms/etc";
+import { maspIndexerUrlAtom, rpcUrlAtom } from "atoms/settings";
 import { queryDependentFn } from "atoms/utils";
 import BigNumber from "bignumber.js";
 import * as osmosis from "chain-registry/mainnet/osmosis";
@@ -21,72 +23,136 @@ import {
   mapNamadaAddressesToAssets,
   mapNamadaAssetsToTokenBalances,
 } from "./functions";
-import { fetchCoinPrices, fetchShieldedBalance } from "./services";
+import {
+  fetchCoinPrices,
+  fetchShieldedBalance,
+  shieldedSync,
+  ShieldedSyncEmitter,
+} from "./services";
 
 export type TokenBalance = AddressWithAsset & {
   amount: BigNumber;
   dollar?: BigNumber;
 };
 
-const DEPRECATED_getViewingKey = async (): Promise<string | undefined> => {
+const DEPRECATED_getViewingKey = async (): Promise<
+  [string, string[]] | undefined
+> => {
   try {
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     const namada: Namada | undefined = (window as any).namada;
     if (namada && semver.lt(namada.version(), "0.3.3")) {
       const accounts = await namada.accounts();
       const defaultAccount = await namada.defaultAccount();
-      const shieldedAccount = accounts?.find(
-        (a) => a.type === "shielded-keys" && a.alias === defaultAccount?.alias
+      const shieldedAccounts = accounts?.filter(
+        (a) => a.type === "shielded-keys"
+      );
+      const defaultShieldedAccount = shieldedAccounts?.find(
+        (a) => a.alias === defaultAccount?.alias
       );
-      return shieldedAccount?.owner;
+      const defaultViewingKey = defaultShieldedAccount?.owner ?? "";
+      const viewingKeys = shieldedAccounts?.map((a) => a.owner ?? "") ?? [];
+
+      return [defaultViewingKey, viewingKeys];
     }
   } catch {
     // do nothing
   }
   return undefined;
 };
 
-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 deprecatedViewingKey = await DEPRECATED_getViewingKey();
       if (deprecatedViewingKey) {
         return deprecatedViewingKey;
       }
-      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 = atomWithQuery<ShieldedSyncEmitter | null>(
+  (get) => {
+    const viewingKeysQuery = get(viewingKeysAtom);
+    const namTokenAddressQuery = get(nativeTokenAddressAtom);
+    const rpcUrl = get(rpcUrlAtom);
+    const maspIndexerUrl = get(maspIndexerUrlAtom);
+
+    return {
+      queryKey: [
+        "shielded-sync",
+        viewingKeysQuery.data,
+        namTokenAddressQuery.data,
+        rpcUrl,
+        maspIndexerUrl,
+      ],
+      ...queryDependentFn(async () => {
+        const viewingKeys = viewingKeysQuery.data;
+        const namTokenAddress = namTokenAddressQuery.data;
+        if (!namTokenAddress || !viewingKeys) {
+          return null;
+        }
+        const [_, allViewingKeys] = viewingKeys;
+        return shieldedSync(
+          rpcUrl,
+          maspIndexerUrl,
+          namTokenAddress,
+          allViewingKeys
+        );
+      }, [viewingKeysQuery, namTokenAddressQuery]),
+    };
+  }
+);
+
 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,
+      shieldedSync.data,
+      rpcUrl,
+      maspIndexerUrl,
     ],
     ...queryDependentFn(async () => {
-      const viewingKey = viewingKeyQuery.data;
+      const viewingKeys = viewingKeysQuery.data;
       const tokenAddresses = tokenAddressesQuery.data;
-      if (!viewingKey || !tokenAddresses) {
+      const syncEmitter = shieldedSync.data;
+      if (!viewingKeys || !tokenAddresses || !syncEmitter) {
         return [];
       }
+      const [viewingKey] = viewingKeys;
+
+      await new Promise<void>((resolve) => {
+        syncEmitter.once(SdkEvents.ProgressBarFinished, () => resolve());
+      });
+
       const response = await fetchShieldedBalance(
         viewingKey,
         tokenAddresses.map((t) => t.address)
@@ -100,7 +166,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,16 @@
-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,
+  ProgressBarFinished,
+  ProgressBarIncremented,
+  ProgressBarStarted,
+  Worker as ShieldedSyncWorkerApi,
+} from "workers/ShieldedSyncWorker";
+import ShieldedSyncWorker from "workers/ShieldedSyncWorker?worker";
 
 const sqsOsmosisApi = "https://sqs.osmosis.zone";
 
@@ -13,6 +24,64 @@ export const fetchCoinPrices = async (
     ).then((res) => res.json())
   : [];
 
+export type ShieldedSyncEventMap = {
+  [SdkEvents.ProgressBarStarted]: ProgressBarStarted[];
+  [SdkEvents.ProgressBarIncremented]: ProgressBarIncremented[];
+  [SdkEvents.ProgressBarFinished]: ProgressBarFinished[];
+};
+
+export type ShieldedSyncEmitter = EventEmitter<ShieldedSyncEventMap>;
+
+let shieldedSyncEmitter: ShieldedSyncEmitter | undefined;
+
+export function shieldedSync(
+  rpcUrl: string,
+  maspIndexerUrl: string,
+  token: string,
+  viewingKeys: string[]
+): EventEmitter<ShieldedSyncEventMap> {
+  if (shieldedSyncEmitter) {
+    return shieldedSyncEmitter;
+  }
+
+  const worker = new ShieldedSyncWorker();
+  const shieldedSyncWorker = Comlink.wrap<ShieldedSyncWorkerApi>(worker);
+  shieldedSyncEmitter = new EventEmitter<ShieldedSyncEventMap>();
+
+  worker.onmessage = (event: MessageEvent<Events>) => {
+    if (!shieldedSyncEmitter) {
+      return;
+    }
+    if (event.data.type === SdkEvents.ProgressBarStarted) {
+      shieldedSyncEmitter.emit(event.data.type, event.data);
+    }
+    if (event.data.type === SdkEvents.ProgressBarIncremented) {
+      shieldedSyncEmitter.emit(event.data.type, event.data);
+    }
+    if (event.data.type === SdkEvents.ProgressBarFinished) {
+      shieldedSyncEmitter.emit(event.data.type, event.data);
+    }
+  };
+
+  (async () => {
+    try {
+      await shieldedSyncWorker.init({
+        type: "init",
+        payload: { rpcUrl, maspIndexerUrl, token },
+      });
+      await shieldedSyncWorker.sync({
+        type: "sync",
+        payload: { vks: viewingKeys },
+      });
+    } finally {
+      worker.terminate();
+      shieldedSyncEmitter = undefined;
+    }
+  })();
+
+  return shieldedSyncEmitter;
+}
+
 export const fetchShieldedBalance = async (
   viewingKey: string,
   addresses: string[]
@@ -21,7 +90,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);
 };