feat: add support for compute budget instructions (#24086)

* Add ComputeBudgetInstruction to web3 sdk

* Prettier fix

* Rename to ComputeBudgetProgram and enable tests

Co-authored-by: Justin Starry <justin@solana.com>

web3.js/src/compute-budget.ts

@@ -0,0 +1,189 @@
+import * as BufferLayout from '@solana/buffer-layout';
+
+import {
+  encodeData,
+  decodeData,
+  InstructionType,
+  IInstructionInputData,
+} from './instruction';
+import {PublicKey} from './publickey';
+import {TransactionInstruction} from './transaction';
+
+/**
+ * Compute Budget Instruction class
+ */
+export class ComputeBudgetInstruction {
+  /**
+   * @internal
+   */
+  constructor() {}
+
+  /**
+   * Decode a compute budget instruction and retrieve the instruction type.
+   */
+  static decodeInstructionType(
+    instruction: TransactionInstruction,
+  ): ComputeBudgetInstructionType {
+    this.checkProgramId(instruction.programId);
+
+    const instructionTypeLayout = BufferLayout.u8('instruction');
+    const typeIndex = instructionTypeLayout.decode(instruction.data);
+
+    let type: ComputeBudgetInstructionType | undefined;
+    for (const [ixType, layout] of Object.entries(
+      COMPUTE_BUDGET_INSTRUCTION_LAYOUTS,
+    )) {
+      if (layout.index == typeIndex) {
+        type = ixType as ComputeBudgetInstructionType;
+        break;
+      }
+    }
+
+    if (!type) {
+      throw new Error(
+        'Instruction type incorrect; not a ComputeBudgetInstruction',
+      );
+    }
+
+    return type;
+  }
+
+  /**
+   * Decode request units compute budget instruction and retrieve the instruction params.
+   */
+  static decodeRequestUnits(
+    instruction: TransactionInstruction,
+  ): RequestUnitsParams {
+    this.checkProgramId(instruction.programId);
+    const {units, additionalFee} = decodeData(
+      COMPUTE_BUDGET_INSTRUCTION_LAYOUTS.RequestUnits,
+      instruction.data,
+    );
+    return {units, additionalFee};
+  }
+
+  /**
+   * Decode request heap frame compute budget instruction and retrieve the instruction params.
+   */
+  static decodeRequestHeapFrame(
+    instruction: TransactionInstruction,
+  ): RequestHeapFrameParams {
+    this.checkProgramId(instruction.programId);
+    const {bytes} = decodeData(
+      COMPUTE_BUDGET_INSTRUCTION_LAYOUTS.RequestHeapFrame,
+      instruction.data,
+    );
+    return {bytes};
+  }
+
+  /**
+   * @internal
+   */
+  static checkProgramId(programId: PublicKey) {
+    if (!programId.equals(ComputeBudgetProgram.programId)) {
+      throw new Error(
+        'invalid instruction; programId is not ComputeBudgetProgram',
+      );
+    }
+  }
+}
+
+/**
+ * An enumeration of valid ComputeBudgetInstructionType's
+ */
+export type ComputeBudgetInstructionType =
+  // FIXME
+  // 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';
+
+type ComputeBudgetInstructionInputData = {
+  RequestUnits: IInstructionInputData & Readonly<RequestUnitsParams>;
+  RequestHeapFrame: IInstructionInputData & Readonly<RequestHeapFrameParams>;
+};
+
+/**
+ * Request units instruction params
+ */
+export interface RequestUnitsParams {
+  /** Units to request for transaction-wide compute */
+  units: number;
+
+  /** Additional fee to pay */
+  additionalFee: number;
+}
+
+/**
+ * Request heap frame instruction params
+ */
+export type RequestHeapFrameParams = {
+  /** Requested transaction-wide program heap size in bytes. Must be multiple of 1024. Applies to each program, including CPIs. */
+  bytes: number;
+};
+
+/**
+ * An enumeration of valid ComputeBudget InstructionType's
+ * @internal
+ */
+export const COMPUTE_BUDGET_INSTRUCTION_LAYOUTS = Object.freeze<{
+  [Instruction in ComputeBudgetInstructionType]: InstructionType<
+    ComputeBudgetInstructionInputData[Instruction]
+  >;
+}>({
+  RequestUnits: {
+    index: 0,
+    layout: BufferLayout.struct<
+      ComputeBudgetInstructionInputData['RequestUnits']
+    >([
+      BufferLayout.u8('instruction'),
+      BufferLayout.u32('units'),
+      BufferLayout.u32('additionalFee'),
+    ]),
+  },
+  RequestHeapFrame: {
+    index: 1,
+    layout: BufferLayout.struct<
+      ComputeBudgetInstructionInputData['RequestHeapFrame']
+    >([BufferLayout.u8('instruction'), BufferLayout.u32('bytes')]),
+  },
+});
+
+/**
+ * Factory class for transaction instructions to interact with the Compute Budget program
+ */
+export class ComputeBudgetProgram {
+  /**
+   * @internal
+   */
+  constructor() {}
+
+  /**
+   * Public key that identifies the Compute Budget program
+   */
+  static programId: PublicKey = new PublicKey(
+    'ComputeBudget111111111111111111111111111111',
+  );
+
+  static requestUnits(params: RequestUnitsParams): TransactionInstruction {
+    const type = COMPUTE_BUDGET_INSTRUCTION_LAYOUTS.RequestUnits;
+    const data = encodeData(type, params);
+    return new TransactionInstruction({
+      keys: [],
+      programId: this.programId,
+      data,
+    });
+  }
+
+  static requestHeapFrame(
+    params: RequestHeapFrameParams,
+  ): TransactionInstruction {
+    const type = COMPUTE_BUDGET_INSTRUCTION_LAYOUTS.RequestHeapFrame;
+    const data = encodeData(type, params);
+    return new TransactionInstruction({
+      keys: [],
+      programId: this.programId,
+      data,
+    });
+  }
+}

web3.js/src/index.ts

@@ -2,6 +2,7 @@ export * from './account';
 export * from './blockhash';
 export * from './bpf-loader-deprecated';
 export * from './bpf-loader';
+export * from './compute-budget';
 export * from './connection';
 export * from './epoch-schedule';
 export * from './ed25519-program';

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

@@ -0,0 +1,186 @@
+import {expect, use} from 'chai';
+import chaiAsPromised from 'chai-as-promised';
+
+import {
+  Keypair,
+  Connection,
+  LAMPORTS_PER_SOL,
+  Transaction,
+  ComputeBudgetProgram,
+  ComputeBudgetInstruction,
+  PublicKey,
+  SystemProgram,
+  sendAndConfirmTransaction,
+} from '../src';
+import {helpers} from './mocks/rpc-http';
+import {url} from './url';
+
+use(chaiAsPromised);
+
+describe('ComputeBudgetProgram', () => {
+  it('requestUnits', () => {
+    const params = {
+      units: 150000,
+      additionalFee: 0,
+    };
+    const transaction = new Transaction().add(
+      ComputeBudgetProgram.requestUnits(params),
+    );
+    expect(transaction.instructions).to.have.length(1);
+    const [computeBudgetInstruction] = transaction.instructions;
+    expect(params).to.eql(
+      ComputeBudgetInstruction.decodeRequestUnits(computeBudgetInstruction),
+    );
+  });
+
+  it('requestHeapFrame', () => {
+    const params = {
+      bytes: 33 * 1024,
+    };
+    const transaction = new Transaction().add(
+      ComputeBudgetProgram.requestHeapFrame(params),
+    );
+    expect(transaction.instructions).to.have.length(1);
+    const [computeBudgetInstruction] = transaction.instructions;
+    expect(params).to.eql(
+      ComputeBudgetInstruction.decodeRequestHeapFrame(computeBudgetInstruction),
+    );
+  });
+
+  it('ComputeBudgetInstruction', () => {
+    const requestUnits = ComputeBudgetProgram.requestUnits({
+      units: 150000,
+      additionalFee: 0,
+    });
+    const requestHeapFrame = ComputeBudgetProgram.requestHeapFrame({
+      bytes: 33 * 1024,
+    });
+
+    const requestUnitsTransaction = new Transaction().add(requestUnits);
+    expect(requestUnitsTransaction.instructions).to.have.length(1);
+    const requestUnitsTransactionType =
+      ComputeBudgetInstruction.decodeInstructionType(
+        requestUnitsTransaction.instructions[0],
+      );
+    expect(requestUnitsTransactionType).to.eq('RequestUnits');
+
+    const requestHeapFrameTransaction = new Transaction().add(requestHeapFrame);
+    expect(requestHeapFrameTransaction.instructions).to.have.length(1);
+    const requestHeapFrameTransactionType =
+      ComputeBudgetInstruction.decodeInstructionType(
+        requestHeapFrameTransaction.instructions[0],
+      );
+    expect(requestHeapFrameTransactionType).to.eq('RequestHeapFrame');
+  });
+
+  if (process.env.TEST_LIVE) {
+    const STARTING_AMOUNT = 2 * LAMPORTS_PER_SOL;
+    const FEE_AMOUNT = LAMPORTS_PER_SOL;
+    it('live compute budget actions', async () => {
+      const connection = new Connection(url, 'confirmed');
+
+      const baseAccount = Keypair.generate();
+      const basePubkey = baseAccount.publicKey;
+      await helpers.airdrop({
+        connection,
+        address: basePubkey,
+        amount: STARTING_AMOUNT,
+      });
+
+      expect(await connection.getBalance(baseAccount.publicKey)).to.eq(
+        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,
+          additionalFee: 2 * FEE_AMOUNT,
+        }),
+        SystemProgram.createAccountWithSeed(createAccountWithSeedParams),
+      );
+      await expect(
+        sendAndConfirmTransaction(
+          connection,
+          createAccountFeeTooHighTransaction,
+          [baseAccount],
+          {preflightCommitment: 'confirmed'},
+        ),
+      ).to.be.rejected;
+
+      expect(await connection.getBalance(baseAccount.publicKey)).to.eq(
+        STARTING_AMOUNT,
+      );
+
+      const createAccountFeeTransaction = new Transaction().add(
+        ComputeBudgetProgram.requestUnits({
+          units: 2,
+          additionalFee: FEE_AMOUNT,
+        }),
+        SystemProgram.createAccountWithSeed(createAccountWithSeedParams),
+      );
+      await sendAndConfirmTransaction(
+        connection,
+        createAccountFeeTransaction,
+        [baseAccount],
+        {preflightCommitment: 'confirmed'},
+      );
+      expect(await connection.getBalance(baseAccount.publicKey)).to.be.at.most(
+        STARTING_AMOUNT - FEE_AMOUNT - minimumAmount,
+      );
+
+      async function expectRequestHeapFailure(bytes: number) {
+        const requestHeapFrameTransaction = new Transaction().add(
+          ComputeBudgetProgram.requestHeapFrame({bytes}),
+        );
+        await expect(
+          sendAndConfirmTransaction(
+            connection,
+            requestHeapFrameTransaction,
+            [baseAccount],
+            {preflightCommitment: 'confirmed'},
+          ),
+        ).to.be.rejected;
+      }
+      const NOT_MULTIPLE_OF_1024 = 33 * 1024 + 1;
+      const BELOW_MIN = 1024;
+      const ABOVE_MAX = 257 * 1024;
+      await expectRequestHeapFailure(NOT_MULTIPLE_OF_1024);
+      await expectRequestHeapFailure(BELOW_MIN);
+      await expectRequestHeapFailure(ABOVE_MAX);
+
+      const VALID_BYTES = 33 * 1024;
+      const requestHeapFrameTransaction = new Transaction().add(
+        ComputeBudgetProgram.requestHeapFrame({bytes: VALID_BYTES}),
+      );
+      await sendAndConfirmTransaction(
+        connection,
+        requestHeapFrameTransaction,
+        [baseAccount],
+        {preflightCommitment: 'confirmed'},
+      );
+    }).timeout(10 * 1000);
+  }
+});