Merge remote-tracking branch 'origin/ph/txWorkerRedis' into ph/txWorkerRedis

Author: arcoraven

Dockerfile

@@ -42,6 +42,7 @@ COPY . .
 # Prune dev dependencies from the packages
 RUN yarn install --frozen-lockfile --production=false --network-timeout 1000000 && \
     yarn build && \
+    yarn copy-files && \
     yarn install --frozen-lockfile --production=true --network-timeout 1000000
 
 ##############################
@@ -64,4 +65,7 @@ COPY --from=build /app/node_modules ./node_modules
 COPY --from=build /app/src/prisma/* ./src/prisma/
 COPY --from=build /app/dist ./dist
 
+# Replace the schema path in the package.json file
+RUN sed -i 's_"schema": "./src/prisma/schema.prisma"_"schema": "./dist/prisma/schema.prisma"_g' package.json
+
 ENTRYPOINT [ "yarn", "start"]

package.json

@@ -22,7 +22,8 @@
     "start:docker-force-build": "docker compose --profile engine --env-file ./.env up --remove-orphans --build",
     "lint:fix": "eslint --fix 'src/**/*.ts'",
     "test:unit": "vitest",
-    "test:coverage": "vitest run --coverage"
+    "test:coverage": "vitest run --coverage",
+    "copy-files": "copyfiles -u 2 ./src/prisma/* ./dist/prisma/ && copyfiles -u 3 ./src/prisma/migrations/**/*.sql ./dist/prisma/migrations/"
   },
   "dependencies": {
     "@aws-sdk/client-kms": "^3.398.0",
@@ -47,6 +48,7 @@
     "bullmq": "^5.11.0",
     "cookie": "^0.5.0",
     "cookie-parser": "^1.4.6",
+    "copyfiles": "^2.4.1",
     "cron-parser": "^4.9.0",
     "crypto": "^1.0.1",
     "crypto-js": "^4.2.0",

src/db/wallets/walletNonce.ts

@@ -1,5 +1,6 @@
 import { Address, eth_getTransactionCount, getRpcClient } from "thirdweb";
 import { getChain } from "../../utils/chain";
+import { logger } from "../../utils/logger";
 import { normalizeAddress } from "../../utils/primitiveTypes";
 import { redis } from "../../utils/redis/redis";
 import { thirdwebClient } from "../../utils/sdk";
@@ -8,16 +9,72 @@ import { thirdwebClient } from "../../utils/sdk";
  * The "last used nonce" stores the last nonce submitted onchain.
  * Example: "25"
  */
-const lastUsedNonceKey = (chainId: number, walletAddress: Address) =>
+export const lastUsedNonceKey = (chainId: number, walletAddress: Address) =>
   `nonce:${chainId}:${normalizeAddress(walletAddress)}`;
 
 /**
- * The "recycled nonces" list stores unsorted nonces to be reused or cancelled.
+ * The "recycled nonces" set stores unsorted nonces to be reused or cancelled.
  * Example: [ "25", "23", "24" ]
  */
-const recycledNoncesKey = (chainId: number, walletAddress: Address) =>
+export const recycledNoncesKey = (chainId: number, walletAddress: Address) =>
   `nonce-recycled:${chainId}:${normalizeAddress(walletAddress)}`;
 
+/**
+ * The "sent nonces" set stores nonces that have been sent on chain but not yet mined.
+ *
+ * Example: [ "25", "23", "24" ]
+ *
+ * The `nonceResyncWorker` periodically fetches the onchain transaction count for each wallet (a),
+ * compares it to the last nonce sent (b), and for every nonce between b and a,
+ * it recycles the nonce if the nonce is not in this set.
+ */
+export const sentNoncesKey = (chainId: number, walletAddress: Address) =>
+  `nonce-sent:${chainId}:${normalizeAddress(walletAddress)}`;
+
+export const splitSentNoncesKey = (key: string) => {
+  const _splittedKeys = key.split(":");
+  const walletAddress = normalizeAddress(_splittedKeys[2]);
+  const chainId = parseInt(_splittedKeys[1]);
+  return { walletAddress, chainId };
+};
+
+/**
+ * Adds a nonce to the sent nonces set (`nonce-sent:${chainId}:${walletAddress}`).
+ */
+export const addSentNonce = async (
+  chainId: number,
+  walletAddress: Address,
+  nonce: number,
+) => {
+  const key = sentNoncesKey(chainId, walletAddress);
+  await redis.sadd(key, nonce.toString());
+};
+
+/**
+ * Removes a nonce from the sent nonces set (`nonce-sent:${chainId}:${walletAddress}`).
+ */
+export const removeSentNonce = async (
+  chainId: number,
+  walletAddress: Address,
+  nonce: number,
+) => {
+  const key = sentNoncesKey(chainId, walletAddress);
+  const removed = await redis.srem(key, nonce.toString());
+  return removed === 1;
+};
+
+/**
+ * Check if a nonce is in the sent nonces set.
+ */
+export const isSentNonce = async (
+  chainId: number,
+  walletAddress: Address,
+  nonce: number,
+) => {
+  const key = sentNoncesKey(chainId, walletAddress);
+  return !!(await redis.sismember(key, nonce.toString()));
+};
+
 /**
  * Acquire an unused nonce.
  * This should be used to send an EOA transaction with this nonce.
@@ -58,8 +115,16 @@ export const recycleNonce = async (
   walletAddress: Address,
   nonce: number,
 ) => {
+  if (isNaN(nonce)) {
+    logger({
+      level: "warn",
+      message: `[recycleNonce] Invalid nonce: ${nonce}`,
+      service: "worker",
+    });
+    return;
+  }
   const key = recycledNoncesKey(chainId, walletAddress);
-  await redis.rpush(key, nonce);
+  await redis.sadd(key, nonce.toString());
 };
 
 /**
@@ -73,7 +138,7 @@ const _acquireRecycledNonce = async (
   walletAddress: Address,
 ) => {
   const key = recycledNoncesKey(chainId, walletAddress);
-  const res = await redis.lpop(key);
+  const res = await redis.spop(key);
   return res ? parseInt(res) : null;
 };
 
@@ -122,3 +187,36 @@ export const deleteAllNonces = async () => {
     await redis.del(keys);
   }
 };
+
+/**
+ * Resync the nonce i.e., max of transactionCount +1 and lastUsedNonce.
+ * @param chainId
+ * @param walletAddress
+ */
+export const rebaseNonce = async (chainId: number, walletAddress: Address) => {
+  const rpcRequest = getRpcClient({
+    client: thirdwebClient,
+    chain: await getChain(chainId),
+  });
+
+  // The next unused nonce = transactionCount.
+  const transactionCount = await eth_getTransactionCount(rpcRequest, {
+    address: walletAddress,
+  });
+
+  // Lua script to set nonce as max
+  const script = `
+    local transactionCount = tonumber(ARGV[1])
+    local lastUsedNonce = tonumber(redis.call('get', KEYS[1]))
+    local nextNonce = math.max(transactionCount-1, lastUsedNonce)
+    redis.call('set', KEYS[1], nextNonce)
+    return nextNonce
+  `;
+  const nextNonce = await redis.eval(
+    script,
+    1,
+    lastUsedNonceKey(chainId, normalizeAddress(walletAddress)),
+    transactionCount.toString(),
+  );
+  return nextNonce;
+};

src/prisma/migrations/20240802175702_transaction_tbl_queued_at_idx/migration.sql

@@ -0,0 +1,2 @@
+-- CreateIndex
+CREATE INDEX CONCURRENTLY "transactions_queuedAt_idx" ON "transactions"("queuedAt");

src/prisma/schema.prisma

@@ -163,6 +163,7 @@ model Transactions {
   @@index([sentAt, minedAt, cancelledAt, errorMessage, queuedAt])
   @@index([sentAt, accountAddress, userOpHash, minedAt, errorMessage, retryCount])
   @@index([sentAt, transactionHash, accountAddress, minedAt, errorMessage, nonce])
+  @@index([queuedAt])
   @@map("transactions")
 }
 

src/scripts/setup-db.ts

@@ -12,7 +12,10 @@ const main = async () => {
       );
     `;
 
-  const schema = `./src/prisma/schema.prisma`;
+  const schema =
+    process.env.NODE_ENV === "production"
+      ? `./dist/prisma/schema.prisma`
+      : `./src/prisma/schema.prisma`;
 
   if (hasWalletsTable) {
     execSync(`yarn prisma migrate reset --force --schema ${schema}`, {

src/server/middleware/adminRoutes.ts

@@ -10,6 +10,7 @@ import { env } from "../../utils/env";
 import { CancelRecycledNoncesQueue } from "../../worker/queues/cancelRecycledNoncesQueue";
 import { MigratePostgresTransactionsQueue } from "../../worker/queues/migratePostgresTransactionsQueue";
 import { MineTransactionQueue } from "../../worker/queues/mineTransactionQueue";
+import { NonceResyncQueue } from "../../worker/queues/nonceResyncQueue";
 import { ProcessEventsLogQueue } from "../../worker/queues/processEventLogsQueue";
 import { ProcessTransactionReceiptsQueue } from "../../worker/queues/processTransactionReceiptsQueue";
 import { PruneTransactionsQueue } from "../../worker/queues/pruneTransactionsQueue";
@@ -28,6 +29,7 @@ const QUEUES: Queue[] = [
   CancelRecycledNoncesQueue.q,
   PruneTransactionsQueue.q,
   MigratePostgresTransactionsQueue.q,
+  NonceResyncQueue.q,
 ];
 
 export const withAdminRoutes = async (fastify: FastifyInstance) => {

src/server/schemas/wallet/index.ts

@@ -5,6 +5,9 @@ export const walletHeaderSchema = Type.Object({
   "x-backend-wallet-address": Type.String({
     examples: ["0x..."],
     description: "Backend wallet address",
+    maxLength: 42,
+    minLength: 42,
+    pattern: "^0x[0-9a-zA-Z]+",
   }),
   "x-idempotency-key": Type.Optional(
     Type.String({
@@ -18,6 +21,9 @@ export const walletWithAAHeaderSchema = Type.Object({
   "x-account-address": Type.Optional(
     Type.String({
       description: "Smart account address",
+      maxLength: 42,
+      minLength: 42,
+      pattern: "^0x[0-9a-zA-Z]+",
     }),
   ),
 });
@@ -34,5 +40,8 @@ export const walletWithAddressParamSchema = Type.Object({
   walletAddress: Type.String({
     examples: ["0x..."],
     description: "Backend wallet address",
+    maxLength: 42,
+    minLength: 42,
+    pattern: "^0x[0-9a-zA-Z]+",
   }),
 });

src/utils/account.ts

@@ -1,5 +1,6 @@
 import { EVMWallet } from "@thirdweb-dev/wallets";
-import { Signer } from "ethers";
+import { Signer, providers } from "ethers";
+
 import { Address } from "thirdweb";
 import { ethers5Adapter } from "thirdweb/adapters/ethers5";
 import { Account } from "thirdweb/wallets";
@@ -9,6 +10,7 @@ import { getAwsKmsWallet } from "../server/utils/wallets/getAwsKmsWallet";
 import { getGcpKmsWallet } from "../server/utils/wallets/getGcpKmsWallet";
 import { getLocalWallet } from "../server/utils/wallets/getLocalWallet";
 import { getSmartWallet } from "../server/utils/wallets/getSmartWallet";
+import { getChain } from "./chain";
 
 export const _accountsCache = new Map<string, Account>();
 
@@ -55,6 +57,7 @@ export const getAccount = async (args: {
 
   // Get smart wallet if `accountAddress` is provided.
   let signer: Signer;
+
   if (accountAddress) {
     const smartWallet = await getSmartWallet({
       chainId,
@@ -66,7 +69,17 @@ export const getAccount = async (args: {
     signer = await wallet.getSigner();
   }
 
-  const account = await ethers5Adapter.signer.fromEthers({ signer });
+  if (walletDetails.type !== WalletType.local) {
+    // Get chain rpc provider.
+    const chain = await getChain(chainId);
+    const provider = new providers.JsonRpcProvider(chain.rpc);
+
+    signer = signer.connect(provider);
+  }
+
+  const account = await ethers5Adapter.signer.fromEthers({
+    signer,
+  });
 
   // Set cache.
   _accountsCache.set(cacheKey, account);

src/worker/index.ts

@@ -10,6 +10,7 @@ import {
 import { initCancelRecycledNoncesWorker } from "./tasks/cancelRecycledNoncesWorker";
 import { initMigratePostgresTransactionsWorker } from "./tasks/migratePostgresTransactionsWorker";
 import { initMineTransactionWorker } from "./tasks/mineTransactionWorker";
+import { initNonceResyncWorker } from "./tasks/nonceResyncWorker";
 import { initProcessEventLogsWorker } from "./tasks/processEventLogsWorker";
 import { initProcessTransactionReceiptsWorker } from "./tasks/processTransactionReceiptsWorker";
 import { initPruneTransactionsWorker } from "./tasks/pruneTransactionsWorker";
@@ -26,6 +27,8 @@ export const initWorker = async () => {
   initSendWebhookWorker();
   await initMigratePostgresTransactionsWorker();
 
+  await initNonceResyncWorker();
+
   // Listen for new & updated configuration data.
   await newConfigurationListener();
   await updatedConfigurationListener();

src/worker/queues/nonceResyncQueue.ts

@@ -0,0 +1,17 @@
+import { Queue } from "bullmq";
+import { redis } from "../../utils/redis/redis";
+import { defaultJobOptions } from "./queues";
+
+export class NonceResyncQueue {
+  static q = new Queue<string>("nonce-resync-cron", {
+    connection: redis,
+    defaultJobOptions,
+  });
+
+  constructor() {
+    NonceResyncQueue.q.setGlobalConcurrency(1);
+
+    // The cron job is defined in `initNonceResyncWorker`
+    // because it requires an async call to query configuration.
+  }
+}

src/worker/tasks/cancelRecycledNoncesWorker.ts

@@ -77,7 +77,7 @@ const getAndDeleteUnusedNonces = async (key: string) => {
   // Returns all unused nonces for this key and deletes the key.
   const script = `
     local key = ARGV[1]
-    local members = redis.call('LRANGE', key, 0, -1)
+    local members = redis.call('SMEMBERS', key)
     redis.call('DEL', key)
     return members
 `;