From 6bbfef7069e7d7a05573fa0ddb25e60ec22b55eb Mon Sep 17 00:00:00 2001 From: Phil Chen Date: Wed, 27 Apr 2022 00:58:32 +0900 Subject: [PATCH] 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 --- web3.js/src/compute-budget.ts | 189 ++++++++++++++++++++++++++++ web3.js/src/index.ts | 1 + web3.js/test/compute-budget.test.ts | 186 +++++++++++++++++++++++++++ 3 files changed, 376 insertions(+) create mode 100644 web3.js/src/compute-budget.ts create mode 100644 web3.js/test/compute-budget.test.ts diff --git a/web3.js/src/compute-budget.ts b/web3.js/src/compute-budget.ts new file mode 100644 index 0000000000..a59721f2cd --- /dev/null +++ b/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; + RequestHeapFrame: IInstructionInputData & Readonly; +}; + +/** + * 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, + }); + } +} diff --git a/web3.js/src/index.ts b/web3.js/src/index.ts index e72a998bbb..efa01b3e07 100644 --- a/web3.js/src/index.ts +++ b/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'; diff --git a/web3.js/test/compute-budget.test.ts b/web3.js/test/compute-budget.test.ts new file mode 100644 index 0000000000..eec2784c45 --- /dev/null +++ b/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); + } +});