feat: support additional compute budget ixs (#25104)

web3.js/src/compute-budget.ts

@@ -1,4 +1,5 @@
 import * as BufferLayout from '@solana/buffer-layout';
+import {u64} from '@solana/buffer-layout-utils';
 
 import {
   encodeData,
@@ -76,6 +77,34 @@ export class ComputeBudgetInstruction {
     return {bytes};
   }
 
+  /**
+   * Decode set compute unit limit compute budget instruction and retrieve the instruction params.
+   */
+  static decodeSetComputeUnitLimit(
+    instruction: TransactionInstruction,
+  ): SetComputeUnitLimitParams {
+    this.checkProgramId(instruction.programId);
+    const {units} = decodeData(
+      COMPUTE_BUDGET_INSTRUCTION_LAYOUTS.SetComputeUnitLimit,
+      instruction.data,
+    );
+    return {units};
+  }
+
+  /**
+   * Decode set compute unit price compute budget instruction and retrieve the instruction params.
+   */
+  static decodeSetComputeUnitPrice(
+    instruction: TransactionInstruction,
+  ): SetComputeUnitPriceParams {
+    this.checkProgramId(instruction.programId);
+    const {microLamports} = decodeData(
+      COMPUTE_BUDGET_INSTRUCTION_LAYOUTS.SetComputeUnitPrice,
+      instruction.data,
+    );
+    return {microLamports};
+  }
+
   /**
    * @internal
    */
@@ -96,11 +125,18 @@ export type ComputeBudgetInstructionType =
   // It would be preferable for this type to be `keyof ComputeBudgetInstructionInputData`
   // but Typedoc does not transpile `keyof` expressions.
   // See https://github.com/TypeStrong/typedoc/issues/1894
-  'RequestUnits' | 'RequestHeapFrame';
+  | 'RequestUnits'
+  | 'RequestHeapFrame'
+  | 'SetComputeUnitLimit'
+  | 'SetComputeUnitPrice';
 
 type ComputeBudgetInstructionInputData = {
   RequestUnits: IInstructionInputData & Readonly<RequestUnitsParams>;
   RequestHeapFrame: IInstructionInputData & Readonly<RequestHeapFrameParams>;
+  SetComputeUnitLimit: IInstructionInputData &
+    Readonly<SetComputeUnitLimitParams>;
+  SetComputeUnitPrice: IInstructionInputData &
+    Readonly<SetComputeUnitPriceParams>;
 };
 
 /**
@@ -109,8 +145,7 @@ type ComputeBudgetInstructionInputData = {
 export interface RequestUnitsParams {
   /** Units to request for transaction-wide compute */
   units: number;
-
+  /** Prioritization fee lamports */
-  /** Additional fee to pay */
   additionalFee: number;
 }
 
@@ -122,6 +157,22 @@ export type RequestHeapFrameParams = {
   bytes: number;
 };
 
+/**
+ * Set compute unit limit instruction params
+ */
+export interface SetComputeUnitLimitParams {
+  /** Transaction-wide compute unit limit */
+  units: number;
+}
+
+/**
+ * Set compute unit price instruction params
+ */
+export interface SetComputeUnitPriceParams {
+  /** Transaction compute unit price used for prioritization fees */
+  microLamports: number | bigint;
+}
+
 /**
  * An enumeration of valid ComputeBudget InstructionType's
  * @internal
@@ -147,6 +198,18 @@ export const COMPUTE_BUDGET_INSTRUCTION_LAYOUTS = Object.freeze<{
       ComputeBudgetInstructionInputData['RequestHeapFrame']
     >([BufferLayout.u8('instruction'), BufferLayout.u32('bytes')]),
   },
+  SetComputeUnitLimit: {
+    index: 2,
+    layout: BufferLayout.struct<
+      ComputeBudgetInstructionInputData['SetComputeUnitLimit']
+    >([BufferLayout.u8('instruction'), BufferLayout.u32('units')]),
+  },
+  SetComputeUnitPrice: {
+    index: 3,
+    layout: BufferLayout.struct<
+      ComputeBudgetInstructionInputData['SetComputeUnitPrice']
+    >([BufferLayout.u8('instruction'), u64('microLamports')]),
+  },
 });
 
 /**
@@ -186,4 +249,30 @@ export class ComputeBudgetProgram {
       data,
     });
   }
+
+  static setComputeUnitLimit(
+    params: SetComputeUnitLimitParams,
+  ): TransactionInstruction {
+    const type = COMPUTE_BUDGET_INSTRUCTION_LAYOUTS.SetComputeUnitLimit;
+    const data = encodeData(type, params);
+    return new TransactionInstruction({
+      keys: [],
+      programId: this.programId,
+      data,
+    });
+  }
+
+  static setComputeUnitPrice(
+    params: SetComputeUnitPriceParams,
+  ): TransactionInstruction {
+    const type = COMPUTE_BUDGET_INSTRUCTION_LAYOUTS.SetComputeUnitPrice;
+    const data = encodeData(type, {
+      microLamports: BigInt(params.microLamports),
+    });
+    return new TransactionInstruction({
+      keys: [],
+      programId: this.programId,
+      data,
+    });
+  }
 }

web3.js/test/compute-budget.test.ts

@@ -8,8 +8,6 @@ import {
   Transaction,
   ComputeBudgetProgram,
   ComputeBudgetInstruction,
-  PublicKey,
-  SystemProgram,
   sendAndConfirmTransaction,
 } from '../src';
 import {helpers} from './mocks/rpc-http';
@@ -21,15 +19,13 @@ describe('ComputeBudgetProgram', () => {
   it('requestUnits', () => {
     const params = {
       units: 150000,
-      additionalFee: 0,
+      additionalFee: LAMPORTS_PER_SOL,
     };
-    const transaction = new Transaction().add(
+    const ix = ComputeBudgetProgram.requestUnits(params);
-      ComputeBudgetProgram.requestUnits(params),
+    const decodedParams = ComputeBudgetInstruction.decodeRequestUnits(ix);
-    );
+    expect(params).to.eql(decodedParams);
-    expect(transaction.instructions).to.have.length(1);
+    expect(ComputeBudgetInstruction.decodeInstructionType(ix)).to.eq(
-    const [computeBudgetInstruction] = transaction.instructions;
+      'RequestUnits',
-    expect(params).to.eql(
-      ComputeBudgetInstruction.decodeRequestUnits(computeBudgetInstruction),
     );
   });
 
@@ -37,48 +33,49 @@ describe('ComputeBudgetProgram', () => {
     const params = {
       bytes: 33 * 1024,
     };
-    const transaction = new Transaction().add(
+    const ix = ComputeBudgetProgram.requestHeapFrame(params);
-      ComputeBudgetProgram.requestHeapFrame(params),
+    const decodedParams = ComputeBudgetInstruction.decodeRequestHeapFrame(ix);
-    );
+    expect(decodedParams).to.eql(params);
-    expect(transaction.instructions).to.have.length(1);
+    expect(ComputeBudgetInstruction.decodeInstructionType(ix)).to.eq(
-    const [computeBudgetInstruction] = transaction.instructions;
+      'RequestHeapFrame',
-    expect(params).to.eql(
-      ComputeBudgetInstruction.decodeRequestHeapFrame(computeBudgetInstruction),
     );
   });
 
-  it('ComputeBudgetInstruction', () => {
+  it('setComputeUnitLimit', () => {
-    const requestUnits = ComputeBudgetProgram.requestUnits({
+    const params = {
-      units: 150000,
+      units: 50_000,
-      additionalFee: 0,
+    };
-    });
+    const ix = ComputeBudgetProgram.setComputeUnitLimit(params);
-    const requestHeapFrame = ComputeBudgetProgram.requestHeapFrame({
+    const decodedParams =
-      bytes: 33 * 1024,
+      ComputeBudgetInstruction.decodeSetComputeUnitLimit(ix);
-    });
+    expect(decodedParams).to.eql(params);
+    expect(ComputeBudgetInstruction.decodeInstructionType(ix)).to.eq(
+      'SetComputeUnitLimit',
+    );
+  });
 
-    const requestUnitsTransaction = new Transaction().add(requestUnits);
+  it('setComputeUnitPrice', () => {
-    expect(requestUnitsTransaction.instructions).to.have.length(1);
+    const params = {
-    const requestUnitsTransactionType =
+      microLamports: 100_000,
-      ComputeBudgetInstruction.decodeInstructionType(
+    };
-        requestUnitsTransaction.instructions[0],
+    const ix = ComputeBudgetProgram.setComputeUnitPrice(params);
-      );
+    const expectedParams = {
-    expect(requestUnitsTransactionType).to.eq('RequestUnits');
+      ...params,
-
+      microLamports: BigInt(params.microLamports),
-    const requestHeapFrameTransaction = new Transaction().add(requestHeapFrame);
+    };
-    expect(requestHeapFrameTransaction.instructions).to.have.length(1);
+    const decodedParams =
-    const requestHeapFrameTransactionType =
+      ComputeBudgetInstruction.decodeSetComputeUnitPrice(ix);
-      ComputeBudgetInstruction.decodeInstructionType(
+    expect(decodedParams).to.eql(expectedParams);
-        requestHeapFrameTransaction.instructions[0],
+    expect(ComputeBudgetInstruction.decodeInstructionType(ix)).to.eq(
-      );
+      'SetComputeUnitPrice',
-    expect(requestHeapFrameTransactionType).to.eq('RequestHeapFrame');
+    );
   });
 
   if (process.env.TEST_LIVE) {
-    const STARTING_AMOUNT = 2 * LAMPORTS_PER_SOL;
+    it('send live request units ix', async () => {
-    const FEE_AMOUNT = LAMPORTS_PER_SOL;
-    it('live compute budget actions', async () => {
       const connection = new Connection(url, 'confirmed');
-
+      const FEE_AMOUNT = LAMPORTS_PER_SOL;
+      const STARTING_AMOUNT = 2 * LAMPORTS_PER_SOL;
       const baseAccount = Keypair.generate();
       const basePubkey = baseAccount.publicKey;
       await helpers.airdrop({
@@ -87,69 +84,49 @@ describe('ComputeBudgetProgram', () => {
         amount: STARTING_AMOUNT,
       });
 
-      expect(await connection.getBalance(baseAccount.publicKey)).to.eq(
+      const additionalFeeTooHighTransaction = new Transaction().add(
-        STARTING_AMOUNT,
-      );
-
-      const seed = 'hi there';
-      const programId = Keypair.generate().publicKey;
-      const createAccountWithSeedAddress = await PublicKey.createWithSeed(
-        basePubkey,
-        seed,
-        programId,
-      );
-      const space = 0;
-
-      let minimumAmount = await connection.getMinimumBalanceForRentExemption(
-        space,
-      );
-
-      const createAccountWithSeedParams = {
-        fromPubkey: basePubkey,
-        newAccountPubkey: createAccountWithSeedAddress,
-        basePubkey,
-        seed,
-        lamports: minimumAmount,
-        space,
-        programId,
-      };
-
-      const createAccountFeeTooHighTransaction = new Transaction().add(
         ComputeBudgetProgram.requestUnits({
-          units: 2,
+          units: 150_000,
-          additionalFee: 2 * FEE_AMOUNT,
+          additionalFee: STARTING_AMOUNT,
         }),
-        SystemProgram.createAccountWithSeed(createAccountWithSeedParams),
       );
+
       await expect(
         sendAndConfirmTransaction(
           connection,
-          createAccountFeeTooHighTransaction,
+          additionalFeeTooHighTransaction,
           [baseAccount],
           {preflightCommitment: 'confirmed'},
         ),
       ).to.be.rejected;
 
-      expect(await connection.getBalance(baseAccount.publicKey)).to.eq(
+      const validAdditionalFeeTransaction = new Transaction().add(
-        STARTING_AMOUNT,
-      );
-
-      const createAccountFeeTransaction = new Transaction().add(
         ComputeBudgetProgram.requestUnits({
-          units: 2,
+          units: 150_000,
           additionalFee: FEE_AMOUNT,
         }),
-        SystemProgram.createAccountWithSeed(createAccountWithSeedParams),
       );
       await sendAndConfirmTransaction(
         connection,
-        createAccountFeeTransaction,
+        validAdditionalFeeTransaction,
         [baseAccount],
         {preflightCommitment: 'confirmed'},
       );
       expect(await connection.getBalance(baseAccount.publicKey)).to.be.at.most(
-        STARTING_AMOUNT - FEE_AMOUNT - minimumAmount,
+        STARTING_AMOUNT - FEE_AMOUNT,
       );
+    });
+
+    it('send live request heap ix', async () => {
+      const connection = new Connection(url, 'confirmed');
+      const STARTING_AMOUNT = 2 * LAMPORTS_PER_SOL;
+      const baseAccount = Keypair.generate();
+      const basePubkey = baseAccount.publicKey;
+      await helpers.airdrop({
+        connection,
+        address: basePubkey,
+        amount: STARTING_AMOUNT,
+      });
 
       async function expectRequestHeapFailure(bytes: number) {
         const requestHeapFrameTransaction = new Transaction().add(
@@ -181,6 +158,63 @@ describe('ComputeBudgetProgram', () => {
         [baseAccount],
         {preflightCommitment: 'confirmed'},
       );
-    }).timeout(10 * 1000);
+    });
+
+    it('send live compute unit ixs', async () => {
+      const connection = new Connection(url, 'confirmed');
+      const FEE_AMOUNT = LAMPORTS_PER_SOL;
+      const STARTING_AMOUNT = 2 * LAMPORTS_PER_SOL;
+      const baseAccount = Keypair.generate();
+      const basePubkey = baseAccount.publicKey;
+      await helpers.airdrop({
+        connection,
+        address: basePubkey,
+        amount: STARTING_AMOUNT,
+      });
+
+      // lamport fee = 2B * 1M / 1M = 2 SOL
+      const prioritizationFeeTooHighTransaction = new Transaction()
+        .add(
+          ComputeBudgetProgram.setComputeUnitPrice({
+            microLamports: 2_000_000_000,
+          }),
+        )
+        .add(
+          ComputeBudgetProgram.setComputeUnitLimit({
+            units: 1_000_000,
+          }),
+        );
+
+      await expect(
+        sendAndConfirmTransaction(
+          connection,
+          prioritizationFeeTooHighTransaction,
+          [baseAccount],
+          {preflightCommitment: 'confirmed'},
+        ),
+      ).to.be.rejected;
+
+      // lamport fee = 1B * 1M / 1M = 1 SOL
+      const validPrioritizationFeeTransaction = new Transaction()
+        .add(
+          ComputeBudgetProgram.setComputeUnitPrice({
+            microLamports: 1_000_000_000,
+          }),
+        )
+        .add(
+          ComputeBudgetProgram.setComputeUnitLimit({
+            units: 1_000_000,
+          }),
+        );
+      await sendAndConfirmTransaction(
+        connection,
+        validPrioritizationFeeTransaction,
+        [baseAccount],
+        {preflightCommitment: 'confirmed'},
+      );
+      expect(await connection.getBalance(baseAccount.publicKey)).to.be.at.most(
+        STARTING_AMOUNT - FEE_AMOUNT,
+      );
+    });
   }
 });