feat(common): add createContract, createNonceManager utils (#1261)

.changeset/tricky-oranges-pump.md

@@ -0,0 +1,10 @@
+---
+"@latticexyz/common": major
+---
+
+Add utils for using viem with MUD
+
+- `createContract` is a wrapper around [viem's `getContract`](https://viem.sh/docs/contract/getContract.html) but with better nonce handling for faster executing of transactions. It has the same arguments and return type as `getContract`.
+- `createNonceManager` helps track local nonces, used by `createContract`.
+
+Also renames `mudTransportObserver` to `transportObserver`.

packages/common/package.json

@@ -61,6 +61,8 @@
     "chalk": "^5.2.0",
     "debug": "^4.3.4",
     "execa": "^7.0.0",
+    "p-queue": "^7.3.4",
+    "p-retry": "^5.1.2",
     "prettier": "^2.8.4",
     "prettier-plugin-solidity": "^1.1.2",
     "viem": "1.3.1"

packages/common/src/createContract.ts

@@ -0,0 +1,125 @@
+import {
+  Abi,
+  Account,
+  Address,
+  Chain,
+  GetContractParameters,
+  GetContractReturnType,
+  Hex,
+  PublicClient,
+  SimulateContractParameters,
+  Transport,
+  WalletClient,
+  WriteContractParameters,
+  getContract,
+} from "viem";
+import pQueue from "p-queue";
+import pRetry from "p-retry";
+import { createNonceManager } from "./createNonceManager";
+import { debug as parentDebug } from "./debug";
+
+const debug = parentDebug.extend("createContract");
+
+// copied from viem because it isn't exported
+// TODO: import from viem?
+function getFunctionParameters(values: [args?: readonly unknown[], options?: object]): {
+  args: readonly unknown[];
+  options: object;
+} {
+  const hasArgs = values.length && Array.isArray(values[0]);
+  const args = hasArgs ? values[0]! : [];
+  const options = (hasArgs ? values[1] : values[0]) ?? {};
+  return { args, options };
+}
+
+export function createContract<
+  TTransport extends Transport,
+  TAddress extends Address,
+  TAbi extends Abi,
+  TChain extends Chain,
+  TAccount extends Account,
+  TPublicClient extends PublicClient<TTransport, TChain>,
+  TWalletClient extends WalletClient<TTransport, TChain, TAccount>
+>({
+  abi,
+  address,
+  publicClient,
+  walletClient,
+}: Required<
+  GetContractParameters<TTransport, TChain, TAccount, TAbi, TPublicClient, TWalletClient, TAddress>
+>): GetContractReturnType<TAbi, TPublicClient, TWalletClient, TAddress> {
+  const contract = getContract<TTransport, TAddress, TAbi, TChain, TAccount, TPublicClient, TWalletClient>({
+    abi,
+    address,
+    publicClient,
+    walletClient,
+  }) as unknown as GetContractReturnType<Abi, PublicClient, WalletClient>;
+
+  if (contract.write) {
+    const nonceManager = createNonceManager({
+      publicClient: publicClient as PublicClient,
+      address: walletClient.account.address,
+    });
+
+    // Concurrency of one means transactions will be queued and inserted into the mem pool synchronously and in order.
+    // Although increasing this will allow for more parallel requests/transactions and nonce errors will get automatically retried,
+    // we can't guarantee local nonce accurancy due to needing async operations (simulate) before incrementing the nonce.
+    const queue = new pQueue({ concurrency: 1 });
+
+    // Replace write calls with our own proxy. Implemented ~the same as viem, but adds better handling of nonces (via queue + retries).
+    contract.write = new Proxy(
+      {},
+      {
+        get(_, functionName: string): GetContractReturnType<Abi, PublicClient, WalletClient>["write"][string] {
+          return async (...parameters) => {
+            const { args, options } = getFunctionParameters(parameters as any);
+
+            async function write(): Promise<Hex> {
+              if (!nonceManager.hasNonce()) {
+                await nonceManager.resetNonce();
+              }
+
+              debug("simulating write", functionName, args, options);
+              const { request } = await publicClient.simulateContract({
+                account: walletClient.account,
+                address,
+                abi,
+                functionName,
+                args,
+                ...options,
+              } as unknown as SimulateContractParameters<TAbi, typeof functionName, TChain>);
+
+              const nonce = nonceManager.nextNonce();
+              debug("calling write function with nonce", nonce, request);
+              const result = await walletClient.writeContract({
+                nonce,
+                ...request,
+              } as unknown as WriteContractParameters<TAbi, typeof functionName, TChain, TAccount>);
+
+              return result;
+            }
+
+            return await queue.add(
+              () =>
+                pRetry(write, {
+                  retries: 3,
+                  onFailedAttempt: async (error) => {
+                    // On nonce errors, reset the nonce and retry
+                    if (nonceManager.shouldResetNonce(error)) {
+                      debug("got nonce error, retrying", error);
+                      await nonceManager.resetNonce();
+                      return;
+                    }
+                    throw error;
+                  },
+                }),
+              { throwOnTimeout: true }
+            );
+          };
+        },
+      }
+    );
+  }
+
+  return contract as unknown as GetContractReturnType<TAbi, TPublicClient, TWalletClient, TAddress>;
+}

packages/common/src/createNonceManager.ts

@@ -0,0 +1,67 @@
+import { BlockTag, Hex, PublicClient } from "viem";
+import { debug as parentDebug } from "./debug";
+
+const debug = parentDebug.extend("createNonceManager");
+
+type CreateNonceManagerOptions = {
+  publicClient: PublicClient;
+  address: Hex;
+  blockTag?: BlockTag;
+};
+
+type CreateNonceManagerResult = {
+  hasNonce: () => boolean;
+  nextNonce: () => number;
+  resetNonce: () => Promise<void>;
+  shouldResetNonce: (error: unknown) => boolean;
+};
+
+export function createNonceManager({
+  publicClient,
+  address,
+  blockTag,
+}: CreateNonceManagerOptions): CreateNonceManagerResult {
+  const nonceRef = { current: -1 };
+  const channel =
+    typeof BroadcastChannel !== "undefined"
+      ? // TODO: fetch chain ID or require it via types?
+        new BroadcastChannel(`mud:createNonceManager:${publicClient.chain?.id}:${address}`)
+      : null;
+
+  if (channel) {
+    channel.addEventListener("message", (event) => {
+      const nonce = JSON.parse(event.data);
+      debug("got nonce from broadcast channel", nonce);
+      nonceRef.current = nonce;
+    });
+  }
+
+  function hasNonce(): boolean {
+    return nonceRef.current >= 0;
+  }
+
+  function nextNonce(): number {
+    if (!hasNonce()) throw new Error("call resetNonce before using nextNonce");
+    const nonce = nonceRef.current++;
+    channel?.postMessage(JSON.stringify(nonceRef.current));
+    return nonce;
+  }
+
+  async function resetNonce(): Promise<void> {
+    const nonce = await publicClient.getTransactionCount({ address, blockTag });
+    nonceRef.current = nonce;
+    channel?.postMessage(JSON.stringify(nonceRef.current));
+    debug("reset nonce to", nonceRef.current);
+  }
+
+  function shouldResetNonce(error: unknown): boolean {
+    return /already known|nonce too low/.test(String(error));
+  }
+
+  return {
+    hasNonce,
+    nextNonce,
+    resetNonce,
+    shouldResetNonce,
+  };
+}

packages/common/src/index.ts

@@ -1,4 +1,6 @@
 export * from "./createBurnerAccount";
+export * from "./createContract";
+export * from "./createNonceManager";
 export * from "./hexToTableId";
-export * from "./mudTransportObserver";
 export * from "./tableIdToHex";
+export * from "./transportObserver";

packages/common/src/mudTransportObserver.ts → packages/common/src/transportObserver.ts

@@ -1,9 +1,9 @@
 import { Hex, Transport, keccak256 } from "viem";
 import { debug as parentDebug } from "./debug";
 
-const debug = parentDebug.extend("mudTransportObserver");
+const debug = parentDebug.extend("transportObserver");
 
-export function mudTransportObserver<TTransport extends Transport>(transport: TTransport): TTransport {
+export function transportObserver<TTransport extends Transport>(transport: TTransport): TTransport {
   return ((opts) => {
     const result = transport(opts);
     const request: typeof result.request = async (req) => {