Set CU limit by simulating transactions (#5)

lorisleiva · web-flow · commit b216b6d14b6b · 2025-02-06T09:11:03.000Z
diff --git a/clients/js/src/createMetadata.ts b/clients/js/src/createMetadata.ts
@@ -80,7 +80,7 @@ export function getCreateMetadataInstructionPlanUsingInstructionData(
     instructions: [
       ...getComputeUnitInstructions({
         computeUnitPrice: input.priorityFees,
-        computeUnitLimit: undefined, // TODO: Add max CU for each instruction.
+        computeUnitLimit: 'simulated',
       }),
       getTransferSolInstruction({
         source: input.payer,
@@ -106,7 +106,7 @@ export function getCreateMetadataInstructionPlanUsingBuffer(
     instructions: [
       ...getComputeUnitInstructions({
         computeUnitPrice: input.priorityFees,
-        computeUnitLimit: undefined, // TODO: Add max CU for each instruction.
+        computeUnitLimit: 'simulated',
       }),
       getTransferSolInstruction({
         source: input.payer,
@@ -145,7 +145,7 @@ export function getCreateMetadataInstructionPlanUsingBuffer(
     instructions: [
       ...getComputeUnitInstructions({
         computeUnitPrice: input.priorityFees,
-        computeUnitLimit: undefined, // TODO: Add max CU for each instruction.
+        computeUnitLimit: 'simulated',
       }),
       getInitializeInstruction({
         ...input,
diff --git a/clients/js/src/instructionPlans/defaultInstructionPlanExecutor.ts b/clients/js/src/instructionPlans/defaultInstructionPlanExecutor.ts
@@ -1,21 +1,57 @@
 import {
-  appendTransactionMessageInstructions,
+  COMPUTE_BUDGET_PROGRAM_ADDRESS,
+  ComputeBudgetInstruction,
+  getSetComputeUnitLimitInstruction,
+  identifyComputeBudgetInstruction,
+} from '@solana-program/compute-budget';
+import {
+  Commitment,
   CompilableTransactionMessage,
+  compileTransaction,
   FullySignedTransaction,
+  getBase64EncodedWireTransaction,
+  GetEpochInfoApi,
+  GetSignatureStatusesApi,
   pipe,
+  Rpc,
+  RpcSubscriptions,
+  sendAndConfirmTransactionFactory,
+  SendTransactionApi,
+  SignatureNotificationsApi,
   signTransactionMessageWithSigners,
+  SimulateTransactionApi,
+  SlotNotificationsApi,
   TransactionMessageWithBlockhashLifetime,
   TransactionMessageWithDurableNonceLifetime,
-  TransactionWithLifetime,
+  TransactionWithBlockhashLifetime,
 } from '@solana/web3.js';
-import { MessageInstructionPlan } from './instructionPlan';
+import {
+  getTransactionMessageFromPlan,
+  MessageInstructionPlan,
+} from './instructionPlan';
 import {
   chunkParallelInstructionPlans,
   createInstructionPlanExecutor,
   InstructionPlanExecutor,
 } from './instructionPlanExecutor';
 
 export type DefaultInstructionPlanExecutorConfig = Readonly<{
+  rpc: Rpc<
+    GetEpochInfoApi &
+      GetSignatureStatusesApi &
+      SendTransactionApi &
+      SimulateTransactionApi
+  >;
+
+  rpcSubscriptions: RpcSubscriptions<
+    SignatureNotificationsApi & SlotNotificationsApi
+  >;
+
+  /**
+   * The commitment to use when confirming transactions.
+   */
+  commitment?: Commitment;
+
   /**
    * When provided, chunks the plans inside a {@link ParallelInstructionPlan}.
    * Each chunk is executed sequentially but each plan within a chunk is
@@ -29,7 +65,7 @@ export type DefaultInstructionPlanExecutorConfig = Readonly<{
    * simulate the transaction to determine the optimal compute unit limit
    * before updating the compute budget instruction with the computed value.
    */
-  simulateComputeUnitLimit?: boolean; // TODO
+  simulateComputeUnitLimit?: boolean;
 
   /**
    * Returns the default transaction message used to send transactions.
@@ -45,33 +81,40 @@ export type DefaultInstructionPlanExecutorConfig = Readonly<{
         | TransactionMessageWithDurableNonceLifetime
       )
   >;
-
-  /**
-   * Sends and confirms a constructed transaction.
-   */
-  sendAndConfirm: (
-    transaction: FullySignedTransaction & TransactionWithLifetime,
-    config?: { abortSignal?: AbortSignal }
-  ) => Promise<void>;
 }>;
 
 export function getDefaultInstructionPlanExecutor(
   config: DefaultInstructionPlanExecutorConfig
 ): InstructionPlanExecutor {
   const {
+    rpc,
+    commitment,
     getDefaultMessage,
     parallelChunkSize: chunkSize,
-    sendAndConfirm,
+    simulateComputeUnitLimit: shouldSimulateComputeUnitLimit,
   } = config;
+  const sendAndConfirm = sendAndConfirmTransactionFactory(config);
 
   return async (plan, config) => {
     const handleMessage = async (plan: MessageInstructionPlan) => {
-      const tx = await pipe(
-        await getDefaultMessage(config),
-        (tx) => appendTransactionMessageInstructions(plan.instructions, tx),
-        (tx) => signTransactionMessageWithSigners(tx)
-      );
-      await sendAndConfirm(tx, config);
+      const defaultMessage = await getDefaultMessage(config);
+      let message = getTransactionMessageFromPlan(defaultMessage, plan);
+
+      if (shouldSimulateComputeUnitLimit) {
+        message = await setComputeUnitLimitBySimulatingTransaction(
+          message,
+          rpc
+        );
+      }
+
+      const tx = (await signTransactionMessageWithSigners(
+        message
+      )) as FullySignedTransaction & TransactionWithBlockhashLifetime;
+      await sendAndConfirm(tx, {
+        ...config,
+        commitment: commitment ?? 'confirmed',
+        skipPreflight: shouldSimulateComputeUnitLimit,
+      });
     };
 
     const executor = pipe(createInstructionPlanExecutor(handleMessage), (e) =>
@@ -81,3 +124,57 @@ export function getDefaultInstructionPlanExecutor(
     return await executor(plan, config);
   };
 }
+
+async function setComputeUnitLimitBySimulatingTransaction<
+  TTransactionMessage extends
+    CompilableTransactionMessage = CompilableTransactionMessage,
+>(
+  message: TTransactionMessage,
+  rpc: Rpc<SimulateTransactionApi>
+): Promise<TTransactionMessage> {
+  const instructionIndex = message.instructions.findIndex((instruction) => {
+    return (
+      instruction.programAddress === COMPUTE_BUDGET_PROGRAM_ADDRESS &&
+      identifyComputeBudgetInstruction(instruction.data as Uint8Array) ===
+        ComputeBudgetInstruction.SetComputeUnitLimit
+    );
+  });
+
+  // Ignore if no compute unit limit instruction is found.
+  if (instructionIndex === -1) {
+    return message;
+  }
+
+  const limit = await getComputeUnitLimitBySimulatingTransaction(message, rpc);
+
+  // Ignore if the limit is not found.
+  if (limit === undefined) {
+    return message;
+  }
+
+  return Object.freeze({
+    ...message,
+    instructions: [
+      ...message.instructions.slice(0, instructionIndex),
+      getSetComputeUnitLimitInstruction({
+        // Use a 1.1x multiplier to the computed limit.
+        units: Number((limit * 110n) / 100n),
+      }),
+      ...message.instructions.slice(instructionIndex + 1),
+    ],
+  } as TTransactionMessage);
+}
+
+async function getComputeUnitLimitBySimulatingTransaction<
+  TTransactionMessage extends
+    CompilableTransactionMessage = CompilableTransactionMessage,
+>(
+  message: TTransactionMessage,
+  rpc: Rpc<SimulateTransactionApi>
+): Promise<bigint | undefined> {
+  const tx = getBase64EncodedWireTransaction(compileTransaction(message));
+  const result = await rpc
+    .simulateTransaction(tx, { encoding: 'base64' })
+    .send();
+  return result.value.unitsConsumed;
+}
diff --git a/clients/js/src/internals.ts b/clients/js/src/internals.ts
@@ -8,7 +8,6 @@ import {
   CompilableTransactionMessage,
   compileTransaction,
   createTransactionMessage,
-  FullySignedTransaction,
   GetAccountInfoApi,
   GetLatestBlockhashApi,
   getTransactionEncoder,
@@ -17,7 +16,6 @@ import {
   pipe,
   ReadonlyUint8Array,
   Rpc,
-  sendAndConfirmTransactionFactory,
   setTransactionMessageFeePayerSigner,
   setTransactionMessageLifetimeUsingBlockhash,
   Transaction,
@@ -129,9 +127,11 @@ function getTimedCacheFunction<T>(
   };
 }
 
+const MAX_COMPUTE_UNIT_LIMIT = 1_400_000;
+
 export function getComputeUnitInstructions(input: {
   computeUnitPrice?: MicroLamports;
-  computeUnitLimit?: number;
+  computeUnitLimit?: number | 'simulated';
 }) {
   const instructions: IInstruction[] = [];
   if (input.computeUnitPrice !== undefined) {
@@ -144,7 +144,10 @@ export function getComputeUnitInstructions(input: {
   if (input.computeUnitLimit !== undefined) {
     instructions.push(
       getSetComputeUnitLimitInstruction({
-        units: input.computeUnitLimit,
+        units:
+          input.computeUnitLimit === 'simulated'
+            ? MAX_COMPUTE_UNIT_LIMIT
+            : input.computeUnitLimit,
       })
     );
   }
@@ -198,7 +201,7 @@ export function getWriteInstructionPlan(input: {
     instructions: [
       ...getComputeUnitInstructions({
         computeUnitPrice: input.priorityFees,
-        computeUnitLimit: undefined, // TODO: Add max CU for each instruction.
+        computeUnitLimit: 'simulated',
       }),
       getWriteInstruction(input),
     ],
@@ -220,15 +223,10 @@ export function getMetadataInstructionPlanExecutor(
   plan: InstructionPlan,
   config?: { abortSignal?: AbortSignal }
 ) => Promise<MetadataResponse> {
-  const sendAndConfirm = sendAndConfirmTransactionFactory(input);
   const executor = getDefaultInstructionPlanExecutor({
+    ...input,
+    simulateComputeUnitLimit: true,
     getDefaultMessage: input.getDefaultMessage,
-    sendAndConfirm: async (tx, config) => {
-      await sendAndConfirm(
-        tx as FullySignedTransaction & TransactionMessageWithBlockhashLifetime,
-        { commitment: input.commitment ?? 'confirmed', ...config }
-      );
-    },
   });
 
   return async (plan, config) => {
diff --git a/clients/js/src/updateMetadata.ts b/clients/js/src/updateMetadata.ts
@@ -116,7 +116,7 @@ export function getUpdateMetadataInstructionPlanUsingInstructionData(
     instructions: [
       ...getComputeUnitInstructions({
         computeUnitPrice: input.priorityFees,
-        computeUnitLimit: undefined, // TODO: Add max CU for each instruction.
+        computeUnitLimit: 'simulated',
       }),
     ],
   };
@@ -164,7 +164,7 @@ export function getUpdateMetadataInstructionPlanUsingBuffer(
     instructions: [
       ...getComputeUnitInstructions({
         computeUnitPrice: input.priorityFees,
-        computeUnitLimit: undefined, // TODO: Add max CU for each instruction.
+        computeUnitLimit: 'simulated',
       }),
     ],
   };
@@ -216,7 +216,7 @@ export function getUpdateMetadataInstructionPlanUsingBuffer(
     instructions: [
       ...getComputeUnitInstructions({
         computeUnitPrice: input.priorityFees,
-        computeUnitLimit: undefined, // TODO: Add max CU for each instruction.
+        computeUnitLimit: 'simulated',
       }),
       getSetDataInstruction({
         ...input,
diff --git a/clients/js/src/utils.ts b/clients/js/src/utils.ts
@@ -21,6 +21,7 @@ import {
   RpcSubscriptions,
   SendTransactionApi,
   SignatureNotificationsApi,
+  SimulateTransactionApi,
   SlotNotificationsApi,
   Transaction,
   TransactionSigner,
@@ -51,6 +52,7 @@ export type MetadataInput = {
       GetEpochInfoApi &
       GetSignatureStatusesApi &
       SendTransactionApi &
+      SimulateTransactionApi &
       GetAccountInfoApi &
       GetMinimumBalanceForRentExemptionApi
   >;
diff --git a/clients/js/test/createMetadata.test.ts b/clients/js/test/createMetadata.test.ts
@@ -55,7 +55,7 @@ test('it creates a canonical metadata account', async (t) => {
   });
 });
 
-test.only('it creates a canonical metadata account with data larger than a transaction size', async (t) => {
+test('it creates a canonical metadata account with data larger than a transaction size', async (t) => {
   // Given the following authority and deployed program.
   const client = createDefaultSolanaClient();
   const authority = await generateKeyPairSignerWithSol(client);
