Improve CU limits for high-frequency transactions

ocp/rpc/transaction/swap_handler.go

@@ -70,7 +70,6 @@ func NewReserveBuySwapHandler(
 		amount:          amount,
 
 		selectedNonce:    selectedNonce,
-		computeUnitLimit: 120_000,
 		computeUnitPrice: 10_000,
 		memoValue:        "buy_v0",
 	}
@@ -152,6 +151,15 @@ func (h *ReserveBuySwapHandler) MakeInstructions(ctx context.Context) ([]solana.
 		return nil, err
 	}
 
+	_, temporaryCoreMintAtaBump, err := token.GetAssociatedAccountAndBump(
+		h.temporaryHolder.PublicKey().ToBytes(),
+		common.CoreMintAccount.PublicKey().ToBytes(),
+	)
+	if err != nil {
+		return nil, err
+	}
+	h.computeUnitLimit = transaction_util.ReserveBuySwapComputeUnitLimit(temporaryCoreMintAtaBump)
+
 	transferFromSourceVmSwapAtaIxn := vm.NewTransferForSwapInstruction(
 		&vm.TransferForSwapInstructionAccounts{
 			VmAuthority: sourceVmConfig.Authority.PublicKey().ToBytes(),
@@ -258,7 +266,6 @@ func NewReserveSellSwapHandler(
 		amount:          amount,
 
 		selectedNonce:    selectedNonce,
-		computeUnitLimit: 145_000,
 		computeUnitPrice: 10_000,
 		memoValue:        "sell_v0",
 	}
@@ -340,6 +347,15 @@ func (h *ReserveSellSwapHandler) MakeInstructions(ctx context.Context) ([]solana
 		return nil, err
 	}
 
+	_, temporarySourceCurrencyAtaBump, err := token.GetAssociatedAccountAndBump(
+		h.temporaryHolder.PublicKey().ToBytes(),
+		h.mint.PublicKey().ToBytes(),
+	)
+	if err != nil {
+		return nil, err
+	}
+	h.computeUnitLimit = transaction_util.ReserveSellSwapComputeUnitLimit(temporarySourceCurrencyAtaBump)
+
 	transferFromSourceVmSwapAtaIxn := vm.NewTransferForSwapInstruction(
 		&vm.TransferForSwapInstructionAccounts{
 			VmAuthority: sourceVmConfig.Authority.PublicKey().ToBytes(),
@@ -449,7 +465,6 @@ func NewReserveBuySellSwapHandler(
 		amount:          amount,
 
 		selectedNonce:    selectedNonce,
-		computeUnitLimit: 250_000,
 		computeUnitPrice: 10_000,
 		memoValue:        "buy_sell_v0",
 	}
@@ -558,6 +573,22 @@ func (h *ReserveBuySellSwapHandler) MakeInstructions(ctx context.Context) ([]sol
 		return nil, err
 	}
 
+	_, temporaryCoreMintAtaBump, err := token.GetAssociatedAccountAndBump(
+		h.temporaryHolder.PublicKey().ToBytes(),
+		common.CoreMintAccount.PublicKey().ToBytes(),
+	)
+	if err != nil {
+		return nil, err
+	}
+	_, temporarySourceCurrencyAtaBump, err := token.GetAssociatedAccountAndBump(
+		h.temporaryHolder.PublicKey().ToBytes(),
+		h.fromMint.PublicKey().ToBytes(),
+	)
+	if err != nil {
+		return nil, err
+	}
+	h.computeUnitLimit = transaction_util.ReserveBuySellSwapComputeUnitLimit(temporaryCoreMintAtaBump, temporarySourceCurrencyAtaBump)
+
 	transferFromSourceVmSwapAtaIxn := vm.NewTransferForSwapInstruction(
 		&vm.TransferForSwapInstructionAccounts{
 			VmAuthority: sourceVmConfig.Authority.PublicKey().ToBytes(),

ocp/transaction/compute_budget.go

@@ -0,0 +1,123 @@
+package transaction
+
+const (
+	// todo: optimize
+	baseExecComputeUnits = 55_000
+
+	// vm check + 4 timelock/vault message derivations
+	numCreateDerivationsInternalExec = 5
+
+	// vm check + 2 message derivations + omnibus invoke_signed
+	numCreateDerivationsExternalExec = 4
+
+	// todo: optimize
+	baseReserveBuySwapComputeUnits     = 80_000 + baseAtaCreateComputeUnits
+	baseReserveSellSwapComputeUnits    = 90_000 + baseAtaCreateComputeUnits
+	baseReserveBuySellSwapComputeUnits = 130_000 + 2*baseAtaCreateComputeUnits
+
+	// todo: optimize
+	baseExternalDepositComputeUnits = 25_000
+	baseCloseVmDepositComputeUnits  = 10_000
+
+	// todo: optimize
+	baseInitTimelockComputeUnits = 20_000 + withdrawReceiptFindComputeUnits
+
+	// init_timelock: vm + memory checks
+	numCreateDerivationsInitTimelock = 2
+
+	// init_timelock derives the withdraw receipt PDA from the VM's PoH value
+	// at execution time, so its bump is unknowable when the transaction is
+	// built. Budget for bump 232, which covers all but ~1 in 16M account
+	// creations.
+	withdrawReceiptFindComputeUnits = 36_000
+
+	// todo: optimize
+	baseAtaCreateComputeUnits = 20_000
+
+	cuPerPdaDerivation = 1_500
+
+	computeUnitMarginPercent = 15
+)
+
+func findPdaComputeUnits(bump uint8) uint32 {
+	return (256 - uint32(bump)) * cuPerPdaDerivation
+}
+
+// WithComputeUnitMargin pads a measured or modeled compute unit count by the
+// standard safety margin.
+func WithComputeUnitMargin(computeUnits uint32) uint32 {
+	return computeUnits * (100 + computeUnitMarginPercent) / 100
+}
+
+// openAccountComputeUnitLimit computes the compute unit limit for an
+// init_timelock transaction. The timelock state, vault and unlock PDAs
+// remain find-derived on-chain even with stored-bump validation, since bump
+// canonicality must be proven at account creation.
+func openAccountComputeUnitLimit(timelockStateBump, vaultBump, unlockBump uint8) uint32 {
+	computeUnits := baseInitTimelockComputeUnits +
+		numCreateDerivationsInitTimelock*cuPerPdaDerivation +
+		findPdaComputeUnits(timelockStateBump) +
+		findPdaComputeUnits(vaultBump) +
+		findPdaComputeUnits(unlockBump)
+	return WithComputeUnitMargin(computeUnits)
+}
+
+func internalExecComputeUnitLimit(numMemoryBanks int) uint32 {
+	return execComputeUnitLimit(baseExecComputeUnits, numCreateDerivationsInternalExec, numMemoryBanks, 0)
+}
+
+func externalExecComputeUnitLimit(numMemoryBanks int) uint32 {
+	return execComputeUnitLimit(baseExecComputeUnits, numCreateDerivationsExternalExec, numMemoryBanks, 0)
+}
+
+func externalExecWithAtaCreateComputeUnitLimit(numMemoryBanks int, ataBump uint8) uint32 {
+	return execComputeUnitLimit(
+		baseExecComputeUnits,
+		numCreateDerivationsExternalExec,
+		numMemoryBanks,
+		baseAtaCreateComputeUnits+findPdaComputeUnits(ataBump),
+	)
+}
+
+func execComputeUnitLimit(baseComputeUnits, numCreateDerivations uint32, numMemoryBanks int, additionalComputeUnits uint32) uint32 {
+	computeUnits := baseComputeUnits + additionalComputeUnits
+	computeUnits += (numCreateDerivations + uint32(numMemoryBanks)) * cuPerPdaDerivation
+	return WithComputeUnitMargin(computeUnits)
+}
+
+// ReserveBuySwapComputeUnitLimit computes the compute unit limit for a
+// reserve buy swap transaction, whose only bump-dependent cost is creating
+// the temporary core mint ATA.
+func ReserveBuySwapComputeUnitLimit(temporaryAtaBump uint8) uint32 {
+	return WithComputeUnitMargin(baseReserveBuySwapComputeUnits + findPdaComputeUnits(temporaryAtaBump))
+}
+
+// ReserveSellSwapComputeUnitLimit computes the compute unit limit for a
+// reserve sell swap transaction, whose only bump-dependent cost is creating
+// the temporary source currency ATA.
+func ReserveSellSwapComputeUnitLimit(temporaryAtaBump uint8) uint32 {
+	return WithComputeUnitMargin(baseReserveSellSwapComputeUnits + findPdaComputeUnits(temporaryAtaBump))
+}
+
+// ReserveBuySellSwapComputeUnitLimit computes the compute unit limit for a
+// reserve buy/sell swap transaction, which creates temporary ATAs for both
+// the core mint and the source currency.
+func ReserveBuySellSwapComputeUnitLimit(temporaryCoreAtaBump, temporarySourceAtaBump uint8) uint32 {
+	return WithComputeUnitMargin(
+		baseReserveBuySellSwapComputeUnits +
+			findPdaComputeUnits(temporaryCoreAtaBump) +
+			findPdaComputeUnits(temporarySourceAtaBump),
+	)
+}
+
+// ExternalDepositComputeUnitLimit computes the compute unit limit for a
+// deposit_from_pda transaction sweeping an external deposit into the VM.
+func ExternalDepositComputeUnitLimit() uint32 {
+	return WithComputeUnitMargin(baseExternalDepositComputeUnits)
+}
+
+// CloseVmDepositComputeUnitLimit computes the compute unit limit for a
+// close_deposit_account_if_empty transaction.
+func CloseVmDepositComputeUnitLimit() uint32 {
+	return WithComputeUnitMargin(baseCloseVmDepositComputeUnits)
+}

ocp/transaction/compute_budget_test.go

@@ -0,0 +1,54 @@
+package transaction
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestExecComputeUnitLimits(t *testing.T) {
+	// base 55,000 + (5 creates + 2 banks) * 1,500 = 65,500, plus 15% margin
+	assert.EqualValues(t, 75_325, internalExecComputeUnitLimit(2))
+
+	// base 55,000 + (4 creates + 2 banks) * 1,500 = 64,000, plus 15% margin
+	assert.EqualValues(t, 73_600, externalExecComputeUnitLimit(2))
+
+	// Adding a memory bank costs one derivation plus margin
+	assert.EqualValues(t, 1_725, internalExecComputeUnitLimit(3)-internalExecComputeUnitLimit(2))
+
+	// The ATA derivation cost scales with how far the bump is from 255
+	canonical := externalExecWithAtaCreateComputeUnitLimit(2, 255)
+	lowBump := externalExecWithAtaCreateComputeUnitLimit(2, 240)
+	assert.EqualValues(t, 15*cuPerPdaDerivation*115/100, lowBump-canonical)
+
+	// A create-on-send transaction always budgets more than a plain external
+	// transfer
+	assert.Greater(t, canonical, externalExecComputeUnitLimit(2))
+}
+
+func TestReserveSwapComputeUnitLimits(t *testing.T) {
+	// base 80,000 + ATA create (20,000) + ATA find (1,500) = 101,500, plus
+	// 15% margin
+	assert.EqualValues(t, 116_725, ReserveBuySwapComputeUnitLimit(255))
+
+	// base 90,000 + ATA create (20,000) + ATA find (1,500) = 111,500, plus
+	// 15% margin
+	assert.EqualValues(t, 128_225, ReserveSellSwapComputeUnitLimit(255))
+
+	// base 130,000 + 2 ATA creates (40,000) + 2 ATA finds (3,000) = 173,000,
+	// plus 15% margin
+	assert.EqualValues(t, 198_950, ReserveBuySellSwapComputeUnitLimit(255, 255))
+}
+
+func TestOpenAccountComputeUnitLimit(t *testing.T) {
+	// base 20,000 + 2 creates (3,000) + state/vault/unlock finds (4,500) +
+	// withdraw receipt allowance (36,000) = 63,500, plus 15% margin
+	assert.EqualValues(t, 73_025, openAccountComputeUnitLimit(255, 255, 255))
+
+	// Low user bumps grow the limit by exactly the extra find iterations
+	assert.EqualValues(
+		t,
+		15*cuPerPdaDerivation*115/100,
+		openAccountComputeUnitLimit(240, 255, 255)-openAccountComputeUnitLimit(255, 255, 255),
+	)
+}

ocp/transaction/transaction.go

@@ -60,7 +60,7 @@ func MakeOpenAccountTransaction(
 
 	instructions := []solana.Instruction{
 		compute_budget.SetComputeUnitPrice(10_000),
-		compute_budget.SetComputeUnitLimit(100_000),
+		compute_budget.SetComputeUnitLimit(openAccountComputeUnitLimit(timelockAccounts.StateBump, timelockAccounts.VaultBump, timelockAccounts.UnlockBump)),
 		initializeInstruction,
 	}
 	return MakeNoncedTransaction(nonce, instructions...)
@@ -148,7 +148,7 @@ func MakeInternalWithdrawTransaction(
 
 	instructions := []solana.Instruction{
 		compute_budget.SetComputeUnitPrice(10_000),
-		compute_budget.SetComputeUnitLimit(100_000),
+		compute_budget.SetComputeUnitLimit(internalExecComputeUnitLimit(mergedMemoryBanks.NumBanks())),
 		execInstruction,
 	}
 	return MakeNoncedTransaction(nonce, instructions...)
@@ -201,7 +201,7 @@ func MakeExternalWithdrawTransaction(
 
 	instructions := []solana.Instruction{
 		compute_budget.SetComputeUnitPrice(10_000),
-		compute_budget.SetComputeUnitLimit(100_000),
+		compute_budget.SetComputeUnitLimit(externalExecComputeUnitLimit(mergedMemoryBanks.NumBanks())),
 		execInstruction,
 	}
 	return MakeNoncedTransaction(nonce, instructions...)
@@ -253,7 +253,7 @@ func MakeInternalTransferWithAuthorityTransaction(
 
 	instructions := []solana.Instruction{
 		compute_budget.SetComputeUnitPrice(10_000),
-		compute_budget.SetComputeUnitLimit(100_000),
+		compute_budget.SetComputeUnitLimit(internalExecComputeUnitLimit(mergedMemoryBanks.NumBanks())),
 		execInstruction,
 	}
 	return MakeNoncedTransaction(nonce, instructions...)
@@ -310,15 +310,7 @@ func MakeExternalTransferWithAuthorityTransaction(
 		},
 	)
 
-	computeLimit := 100_000
-	if isCreateOnSend {
-		computeLimit = 125_000
-	}
-
-	instructions := []solana.Instruction{
-		compute_budget.SetComputeUnitPrice(10_000),
-		compute_budget.SetComputeUnitLimit(uint32(computeLimit)),
-	}
+	var instructions []solana.Instruction
 	if isCreateOnSend {
 		if externalDestinationOwner == nil {
 			return solana.Transaction{}, errors.New("destination owner is required")
@@ -335,7 +327,24 @@ func MakeExternalTransferWithAuthorityTransaction(
 			return solana.Transaction{}, errors.New("invalid destination owner")
 		}
 
-		instructions = append(instructions, createIdempotentInstruction)
+		_, ataBump, err := token.GetAssociatedAccountAndBump(
+			externalDestinationOwner.PublicKey().ToBytes(),
+			mint.PublicKey().ToBytes(),
+		)
+		if err != nil {
+			return solana.Transaction{}, err
+		}
+
+		instructions = []solana.Instruction{
+			compute_budget.SetComputeUnitPrice(10_000),
+			compute_budget.SetComputeUnitLimit(externalExecWithAtaCreateComputeUnitLimit(mergedMemoryBanks.NumBanks(), ataBump)),
+			createIdempotentInstruction,
+		}
+	} else {
+		instructions = []solana.Instruction{
+			compute_budget.SetComputeUnitPrice(10_000),
+			compute_budget.SetComputeUnitLimit(externalExecComputeUnitLimit(mergedMemoryBanks.NumBanks())),
+		}
 	}
 	instructions = append(instructions, execInstruction)
 	return MakeNoncedTransaction(nonce, instructions...)
@@ -400,3 +409,13 @@ func MergeMemoryBanks(accounts ...*common.Account) (*MergedMemoryBankResult, err
 		Indices: indices,
 	}, nil
 }
+
+func (r *MergedMemoryBankResult) NumBanks() int {
+	var count int
+	for _, bank := range []*ed25519.PublicKey{r.A, r.B, r.C, r.D} {
+		if bank != nil {
+			count++
+		}
+	}
+	return count
+}

ocp/worker/currency/feeburner/worker.go

@@ -17,8 +17,10 @@ import (
 )
 
 const (
-	burnFeesComputeUnitLimit = 100_000
-	computeUnitPrice         = 10_000
+	perBurnComputeUnits = 5_000
+	baseComputeUnits    = 300
+
+	computeUnitPrice = 10_000
 )
 
 type burnTarget struct {
@@ -161,7 +163,7 @@ func (p *runtime) makeBurnTransaction(batch []*burnTarget) solana.Transaction {
 	ixns := make([]solana.Instruction, 0, len(batch)+2)
 	ixns = append(
 		ixns,
-		compute_budget.SetComputeUnitLimit(uint32(len(batch))*burnFeesComputeUnitLimit),
+		compute_budget.SetComputeUnitLimit(transaction_util.WithComputeUnitMargin(baseComputeUnits+uint32(len(batch))*perBurnComputeUnits)),
 		compute_budget.SetComputeUnitPrice(computeUnitPrice),
 	)
 	for _, target := range batch {

ocp/worker/currency/launcher/util.go

@@ -802,7 +802,7 @@ func (p *runtime) populateNonceMemory(ctx context.Context, accounts *newCurrency
 
 			err := func() error {
 				ixns := []solana.Instruction{
-					compute_budget.SetComputeUnitLimit(550_000),
+					compute_budget.SetComputeUnitLimit(400_000),
 					compute_budget.SetComputeUnitPrice(10_000),
 				}
 				for i := range initVdnIxnsPerTxn {

ocp/worker/geyser/external_deposit.go

@@ -120,7 +120,7 @@ func initiateExternalDepositIntoVm(ctx context.Context, data ocp_data.Provider,
 		vmConfig.Authority.PublicKey().ToBytes(),
 		memov2.Instruction(codeVmDepositMemoValue),
 		compute_budget.SetComputeUnitPrice(10_000),
-		compute_budget.SetComputeUnitLimit(45_000),
+		compute_budget.SetComputeUnitLimit(transaction_util.ExternalDepositComputeUnitLimit()),
 		vm.NewDepositFromPdaInstruction(
 			&vm.DepositFromPdaInstructionAccounts{
 				VmAuthority: vmConfig.Authority.PublicKey().ToBytes(),
@@ -437,7 +437,7 @@ func closeVmDepositAccount(ctx context.Context, data ocp_data.Provider, userAuth
 	txn := solana.NewLegacyTransaction(
 		vmConfig.Authority.PublicKey().ToBytes(),
 		compute_budget.SetComputeUnitPrice(10_000),
-		compute_budget.SetComputeUnitLimit(25_000),
+		compute_budget.SetComputeUnitLimit(transaction_util.CloseVmDepositComputeUnitLimit()),
 		vm.NewCloseDepositAccountIfEmptyInstruction(
 			&vm.CloseDepositAccountIfEmptyInstructionAccounts{
 				VmAuthority: vmConfig.Authority.PublicKey().ToBytes(),