Add transaction locking for nonce management (#132)

* Update to make transaction processing a prisma.

* Update queries to support pgtx

* Use transaction with timeout

* Add locking on getWalletNonce

* Update nodemon to watch src

Author: adam-maj

core/env.ts

@@ -61,7 +61,7 @@ export const env = createEnv({
     MAX_PRIORITY_FEE_PER_GAS_FOR_RETRY: z.string().default("55000000000"),
     MAX_RETRIES_FOR_TX: z.coerce.number().default(3),
     RETRY_TX_CRON_SCHEDULE: z.string().default("*/30 * * * * *"),
-    MAX_BLOCKS_ELAPSED_BEFORE_RETRY: z.coerce.number().default(50),
+    MAX_BLOCKS_ELAPSED_BEFORE_RETRY: z.coerce.number().default(10),
     MAX_WAIT_TIME_BEFORE_RETRY: z.coerce.number().default(600),
   },
   clientPrefix: "NEVER_USED",

core/sdk/sdk.ts

@@ -11,6 +11,7 @@ import {
   walletTableSchema,
 } from "../../server/schemas/wallet";
 import { getWalletDetails } from "../../src/db/wallets/getWalletDetails";
+import { PrismaTransaction } from "../../src/schema/prisma";
 import { env } from "../env";
 import { networkResponseSchema } from "../schema";
 
@@ -95,7 +96,11 @@ const cacheSdk = (
 };
 
 const walletDataMap: Map<string, string> = new Map();
-const getCachedWallet = async (walletAddress: string, chainId: number) => {
+const getCachedWallet = async (
+  walletAddress: string,
+  chainId: number,
+  pgtx?: PrismaTransaction,
+) => {
   walletAddress = walletAddress.toLowerCase();
   let walletData;
   const cachedWallet = walletDataMap.get(walletAddress);
@@ -104,7 +109,11 @@ const getCachedWallet = async (walletAddress: string, chainId: number) => {
   } else {
     console.log("Checking details for address", walletAddress);
     // TODO: This needs to be changed...
-    walletData = await getWalletDetails({ address: walletAddress, chainId });
+    walletData = await getWalletDetails({
+      pgtx,
+      address: walletAddress,
+      chainId,
+    });
     console.log("Received wallet data:", walletData);
     if (walletData) {
       walletDataMap.set(walletAddress, JSON.stringify(walletData));
@@ -121,6 +130,7 @@ const THIRDWEB_API_SECRET_KEY = env.THIRDWEB_API_SECRET_KEY;
 export const getSDK = async (
   chainName: ChainOrRpc,
   walletAddress?: string,
+  pgtx?: PrismaTransaction,
 ): Promise<ThirdwebSDK> => {
   let walletData: Static<typeof walletTableSchema> | undefined;
 
@@ -158,7 +168,7 @@ export const getSDK = async (
     return sdk;
   }
 
-  walletData = await getCachedWallet(walletAddress, chain.chainId);
+  walletData = await getCachedWallet(walletAddress, chain.chainId, pgtx);
 
   if (!walletData) {
     throw new Error(`Wallet not found for address: ${walletAddress}`);

package.json

@@ -10,8 +10,8 @@
     "docker": "docker compose --env-file ./.env up --remove-orphans",
     "docker:build": "docker compose build --no-cache",
     "dev": "yarn dev:infra && yarn prisma:init && yarn dev:server & sleep 10 && yarn dev:worker",
-    "dev:server": "yarn prisma:init && nodemon --watch 'server/**/*.ts' --watch 'core/**/*.ts' --exec 'npx tsx ./server/index.ts' --files server/index.ts",
-    "dev:worker": "yarn prisma:init && nodemon --watch 'worker/**/*.ts' --watch 'core/**/*.ts' --exec 'npx tsx ./worker/index.ts' --files worker/index.ts",
+    "dev:server": "yarn prisma:init && nodemon --watch 'server/**/*.ts' --watch 'core/**/*.ts' --watch 'src/**/*.ts' --exec 'npx tsx ./server/index.ts' --files server/index.ts",
+    "dev:worker": "yarn prisma:init && nodemon --watch 'worker/**/*.ts' --watch 'core/**/*.ts' --watch 'src/**/*.ts' --exec 'npx tsx ./worker/index.ts' --files worker/index.ts",
     "dev:infra": "docker compose -f ./docker-compose-infra.yml up -d",
     "build": "yarn && rm -rf dist && tsc -p ./tsconfig.json --outDir dist",
     "prisma:reset": "prisma migrate reset --force && prisma generate",

src/db/client.ts

@@ -1,3 +1,8 @@
 import { PrismaClient } from "@prisma/client";
+import { PrismaTransaction } from "../schema/prisma";
 
 export const prisma = new PrismaClient();
+
+export const getPrismaWithPostgresTx = (pgtx?: PrismaTransaction) => {
+  return pgtx || prisma;
+};

src/db/transactions/getAllTxs.ts

@@ -4,17 +4,20 @@ import {
   transactionResponseSchema,
 } from "../../../server/schemas/transaction";
 import { ContractExtension } from "../../schema/extension";
-import { prisma } from "../client";
+import { PrismaTransaction } from "../../schema/prisma";
+import { getPrismaWithPostgresTx } from "../client";
 import { cleanTxs } from "./cleanTxs";
 
 interface GetAllTxsParams {
+  pgtx?: PrismaTransaction;
   page: number;
   limit: number;
   filter?: TransactionStatusEnum;
   extensions?: ContractExtension[];
 }
 
 export const getAllTxs = async ({
+  pgtx,
   page,
   limit,
   filter,
@@ -40,6 +43,8 @@ export const getAllTxs = async ({
     filterBy = "errorMessage";
   }
 
+  const prisma = getPrismaWithPostgresTx(pgtx);
+
   // TODO: Cleaning should be handled by zod
   const txs = await prisma.transactions.findMany({
     where: {

src/db/transactions/getQueuedTxs.ts

@@ -1,12 +1,19 @@
 import { Static } from "@sinclair/typebox";
 import { env } from "../../../core/env";
 import { transactionResponseSchema } from "../../../server/schemas/transaction";
-import { prisma } from "../client";
+import { PrismaTransaction } from "../../schema/prisma";
+import { getPrismaWithPostgresTx } from "../client";
 import { cleanTxs } from "./cleanTxs";
 
-export const getQueuedTxs = async (): Promise<
+interface GetQueuedTxsParams {
+  pgtx?: PrismaTransaction;
+}
+
+export const getQueuedTxs = async ({ pgtx }: GetQueuedTxsParams = {}): Promise<
   Static<typeof transactionResponseSchema>[]
 > => {
+  const prisma = getPrismaWithPostgresTx(pgtx);
+
   // TODO: Don't use env var for transactions to batch
   const txs = await prisma.$queryRaw`
 SELECT

src/db/transactions/getSentTxs.ts

@@ -1,12 +1,19 @@
 import { Static } from "@sinclair/typebox";
 import { env } from "../../../core/env";
 import { transactionResponseSchema } from "../../../server/schemas/transaction";
-import { prisma } from "../client";
+import { PrismaTransaction } from "../../schema/prisma";
+import { getPrismaWithPostgresTx } from "../client";
 import { cleanTxs } from "./cleanTxs";
 
-export const getSentTxs = async (): Promise<
+interface GetSentTxsParams {
+  pgtx?: PrismaTransaction;
+}
+
+export const getSentTxs = async ({ pgtx }: GetSentTxsParams = {}): Promise<
   Static<typeof transactionResponseSchema>[]
 > => {
+  const prisma = getPrismaWithPostgresTx(pgtx);
+
   const txs = await prisma.transactions.findMany({
     where: {
       processedAt: {

src/db/transactions/getTxById.ts

@@ -1,11 +1,15 @@
-import { prisma } from "../client";
+import { PrismaTransaction } from "../../schema/prisma";
+import { getPrismaWithPostgresTx } from "../client";
 import { cleanTxs } from "./cleanTxs";
 
 interface GetTxByIdParams {
+  pgtx?: PrismaTransaction;
   queueId: string;
 }
 
-export const getTxById = async ({ queueId }: GetTxByIdParams) => {
+export const getTxById = async ({ pgtx, queueId }: GetTxByIdParams) => {
+  const prisma = getPrismaWithPostgresTx(pgtx);
+
   const tx = await prisma.transactions.findUnique({
     where: {
       id: queueId,

src/db/transactions/getTxToRetry.ts

@@ -1,11 +1,18 @@
 import { Static } from "@sinclair/typebox";
 import { transactionResponseSchema } from "../../../server/schemas/transaction";
-import { prisma } from "../client";
+import type { PrismaTransaction } from "../../schema/prisma";
+import { getPrismaWithPostgresTx } from "../client";
 import { cleanTxs } from "./cleanTxs";
 
-export const getTxToRetry = async (): Promise<
+interface GetTxToRetryParams {
+  pgtx?: PrismaTransaction;
+}
+
+export const getTxToRetry = async ({ pgtx }: GetTxToRetryParams = {}): Promise<
   Static<typeof transactionResponseSchema>[]
 > => {
+  const prisma = getPrismaWithPostgresTx(pgtx);
+
   // TODO: Why is this checking that transaction hash is not null
   const txs = await prisma.$queryRaw`
 SELECT

src/db/transactions/queueTx.ts

@@ -1,9 +1,11 @@
 import type { DeployTransaction, Transaction } from "@thirdweb-dev/sdk";
 import { BigNumber } from "ethers";
 import type { ContractExtension } from "../../schema/extension";
-import { prisma } from "../client";
+import { PrismaTransaction } from "../../schema/prisma";
+import { getPrismaWithPostgresTx } from "../client";
 
 interface QueueTxParams {
+  pgtx?: PrismaTransaction;
   tx: Transaction<any> | DeployTransaction;
   chainId: number;
   extension: ContractExtension;
@@ -14,12 +16,15 @@ interface QueueTxParams {
 
 // TODO: Simulation should be done before this function...
 export const queueTx = async ({
+  pgtx,
   tx,
   chainId,
   extension,
   deployedContractAddress,
   deployedContractType,
 }: QueueTxParams) => {
+  const prisma = getPrismaWithPostgresTx(pgtx);
+
   // TODO: SDK should have a JSON.stringify() method.
   const fromAddress = (await tx.getSignerAddress()).toLowerCase();
   const toAddress = tx.getTarget().toLowerCase();

src/db/transactions/retryTx.ts

@@ -1,18 +1,21 @@
-import { prisma } from "../client";
+import { PrismaTransaction } from "../../schema/prisma";
+import { getPrismaWithPostgresTx } from "../client";
 
 interface RetryTxParams {
+  pgtx?: PrismaTransaction;
   queueId: string;
   maxFeePerGas: string;
   maxPriorityFeePerGas: string;
 }
 
-// TODO: Switch all functions to object params
 export const retryTx = async ({
+  pgtx,
   queueId,
   maxFeePerGas,
   maxPriorityFeePerGas,
 }: RetryTxParams) => {
-  console.log("Retrying tx...");
+  const prisma = getPrismaWithPostgresTx(pgtx);
+
   await prisma.transactions.update({
     where: {
       id: queueId,

src/db/transactions/updateTx.ts

@@ -1,9 +1,11 @@
 import { Transactions } from "@prisma/client";
 import { providers } from "ethers";
 import { TransactionStatusEnum } from "../../../server/schemas/transaction";
-import { prisma } from "../client";
+import { PrismaTransaction } from "../../schema/prisma";
+import { getPrismaWithPostgresTx } from "../client";
 
 interface UpdateTxParams {
+  pgtx?: PrismaTransaction;
   queueId: string;
   status: TransactionStatusEnum;
   // TODO: Receipt never actually gets used here... should get passed through.
@@ -12,11 +14,14 @@ interface UpdateTxParams {
 }
 
 export const updateTx = async ({
+  pgtx,
   queueId,
   status,
   res,
   txData,
 }: UpdateTxParams) => {
+  const prisma = getPrismaWithPostgresTx(pgtx);
+
   switch (status) {
     case TransactionStatusEnum.Submitted:
       await prisma.transactions.update({

src/db/wallets/cleanWallet.ts

@@ -1,6 +1,6 @@
 import { WalletDetails, WalletNonce } from "@prisma/client";
 
-type WalletNonceWithDetails = WalletNonce & {
+export type WalletNonceWithDetails = WalletNonce & {
   walletDetails: WalletDetails;
 };
 

src/db/wallets/createWalletDetails.ts

@@ -1,8 +1,10 @@
+import { PrismaTransaction } from "../../schema/prisma";
 import type { WalletType } from "../../schema/wallet";
-import { prisma } from "../client";
+import { getPrismaWithPostgresTx } from "../client";
 
 // TODO: Case on types by wallet type
 interface CreateWalletDetailsParams {
+  pgtx?: PrismaTransaction;
   address: string;
   type: WalletType;
   awsKmsKeyId?: string;
@@ -14,9 +16,12 @@ interface CreateWalletDetailsParams {
   gcpKmsResourcePath?: string;
 }
 
-export const createWalletDetails = async (
-  walletDetails: CreateWalletDetailsParams,
-) => {
+export const createWalletDetails = async ({
+  pgtx,
+  ...walletDetails
+}: CreateWalletDetailsParams) => {
+  const prisma = getPrismaWithPostgresTx(pgtx);
+
   return prisma.walletDetails.create({
     data: {
       ...walletDetails,

src/db/wallets/createWalletNonce.ts

@@ -1,17 +1,22 @@
 import { BigNumber } from "ethers";
 import { getSDK } from "../../../core/sdk/sdk";
 import { getWalletNonce } from "../../../core/services/blockchain";
-import { prisma } from "../client";
+import { PrismaTransaction } from "../../schema/prisma";
+import { getPrismaWithPostgresTx } from "../client";
 
 interface CreateWalletNonceParams {
+  pgtx?: PrismaTransaction;
   chainId: number;
   address: string;
 }
 
 export const createWalletNonce = async ({
+  pgtx,
   chainId,
   address,
 }: CreateWalletNonceParams) => {
+  const prisma = getPrismaWithPostgresTx(pgtx);
+
   // TODO: chainId instead of chainName being passed around everywhere
   // or just pass SDK around
   const sdk = await getSDK(chainId.toString());

src/db/wallets/getAllWallets.ts

@@ -1,19 +1,31 @@
-import { prisma } from "../client";
+import { PrismaTransaction } from "../../schema/prisma";
+import { getPrismaWithPostgresTx } from "../client";
 import { cleanWallet } from "../wallets/cleanWallet";
 
-export interface GetAllWalletsParams {
-  chainId: number;
+interface GetAllWalletsParams {
+  pgtx?: PrismaTransaction;
 }
 
 // TODO: Add error logging handler from req.log to all queries
-export const getAllWallets = async () => {
+export const getAllWallets = async ({ pgtx }: GetAllWalletsParams = {}) => {
+  const prisma = getPrismaWithPostgresTx(pgtx);
+
   const wallets = await prisma.walletDetails.findMany();
   return wallets;
 };
 
+interface GetAllWalletsByChainIdParams {
+  pgtx: PrismaTransaction;
+  chainId: number;
+}
+
+// TODO: Change to chainId
 export const getAllWalletsByChain = async ({
+  pgtx,
   chainId,
-}: GetAllWalletsParams) => {
+}: GetAllWalletsByChainIdParams) => {
+  const prisma = getPrismaWithPostgresTx(pgtx);
+
   const wallets = await prisma.walletNonce.findMany({
     where: {
       chainId,