implement copy-on-write account handling to cut down on allocations

pkg/fees/fees.go

@@ -139,8 +139,11 @@ func CalculateAndDeductTxFees(tx *solana.Transaction, txMeta *rpc.TransactionMet
 	}
 	////mlog.Log.Debugf("feePayerAcct.Lamports=%d totalTxFee=%d", feePayerAcct.Lamports, totalTxFee)
 
+	feePayerAcct, err = transactionAccts.Touch(feePayerIdx)
+	if err != nil {
+		return feeInfo, 0, err
+	}
 	feePayerAcct.Lamports -= totalTxFee
-	transactionAccts.Touch(feePayerIdx)
 
 	return feeInfo, feePayerAcct.Lamports, nil
 }

pkg/rent/rent.go

@@ -113,8 +113,11 @@ func VerifyRentStateChanges(preStates []*RentStateInfo, postStates []*RentStateI
 func MaybeSetRentExemptRentEpochMax(slotCtx *sealevel.SlotCtx, rent *sealevel.SysvarRent, f *features.Features, txAccts *sealevel.TransactionAccounts) {
 	for idx := range txAccts.Accounts {
 		if ShouldSetRentExemptRentEpochMax(slotCtx, rent, f, txAccts.Accounts[idx]) {
-			txAccts.Accounts[idx].RentEpoch = math.MaxUint64
-			txAccts.Touch(uint64(idx))
+			touchedAcct, err := txAccts.Touch(uint64(idx))
+			if err != nil {
+				panic("unable to mark rent-exempt account as touched")
+			}
+			touchedAcct.RentEpoch = math.MaxUint64
 		}
 	}
 }

pkg/replay/accounts.go

@@ -14,9 +14,13 @@ import (
 
 // Account clone tracking for profiling copy-on-write optimization potential
 var (
-	// Per-transaction account clone stats (loaded in loadAndValidateTxAcctsSimd186)
-	TxAcctsCloned      atomic.Uint64 // Total accounts cloned across all txs
-	TxAcctsClonedBytes atomic.Uint64 // Total bytes of account data cloned
+	// Per-transaction account load stats (accounts referenced by tx execution)
+	TxAcctsLoaded      atomic.Uint64 // Total accounts loaded into tx contexts
+	TxAcctsLoadedBytes atomic.Uint64 // Total bytes referenced by tx contexts
+
+	// Per-transaction copy-on-write clone stats (first write in TransactionAccounts.Touch)
+	TxAcctsCloned      atomic.Uint64 // Total accounts cloned on first write
+	TxAcctsClonedBytes atomic.Uint64 // Total bytes cloned on first write
 
 	// Per-transaction modification stats (touched in handleModifiedAccounts)
 	TxAcctsTouched      atomic.Uint64 // Total accounts actually modified
@@ -28,24 +32,33 @@ var (
 
 // CloneStats holds account clone/modify metrics for reporting
 type CloneStats struct {
-	AcctsCloned      uint64 // Accounts loaded (cloned)
-	AcctsClonedBytes uint64 // Bytes cloned
-	AcctsTouched     uint64 // Accounts modified
+	AcctsLoaded       uint64 // Accounts loaded into tx contexts
+	AcctsLoadedBytes  uint64 // Bytes referenced by tx contexts
+	AcctsCloned       uint64 // Accounts loaded (cloned)
+	AcctsClonedBytes  uint64 // Bytes cloned
+	AcctsTouched      uint64 // Accounts modified
 	AcctsTouchedBytes uint64 // Bytes of modified accounts
-	TxCount          uint64 // Number of transactions
+	TxCount           uint64 // Number of transactions
 }
 
 // GetAndResetCloneStats returns current clone stats and resets counters
 func GetAndResetCloneStats() CloneStats {
 	return CloneStats{
-		AcctsCloned:      TxAcctsCloned.Swap(0),
-		AcctsClonedBytes: TxAcctsClonedBytes.Swap(0),
-		AcctsTouched:     TxAcctsTouched.Swap(0),
+		AcctsLoaded:       TxAcctsLoaded.Swap(0),
+		AcctsLoadedBytes:  TxAcctsLoadedBytes.Swap(0),
+		AcctsCloned:       TxAcctsCloned.Swap(0),
+		AcctsClonedBytes:  TxAcctsClonedBytes.Swap(0),
+		AcctsTouched:      TxAcctsTouched.Swap(0),
 		AcctsTouchedBytes: TxAcctsTouchedBytes.Swap(0),
-		TxCount:          TxCount.Swap(0),
+		TxCount:           TxCount.Swap(0),
 	}
 }
 
+func recordTxAcctCowClone(acct *accounts.Account) {
+	TxAcctsCloned.Add(1)
+	TxAcctsClonedBytes.Add(uint64(len(acct.Data)))
+}
+
 func loadAndValidateTxAccts(slotCtx *sealevel.SlotCtx, acctMetasPerInstr [][]sealevel.AccountMeta, tx *solana.Transaction, instrs []sealevel.Instruction, instrsAcct *accounts.Account, loadedAcctBytesLimit uint32) (*sealevel.TransactionAccounts, []*solana.AccountMeta, error) {
 	txAcctMetas, err := tx.AccountMetaList()
 	if err != nil {
@@ -63,29 +76,34 @@ func loadAndValidateTxAccts(slotCtx *sealevel.SlotCtx, acctMetasPerInstr [][]sea
 		}
 	}
 
-	acctsForTx := make([]accounts.Account, 0, len(txAcctMetas))
+	acctsForTx := make([]*accounts.Account, 0, len(txAcctMetas))
+	acctsShared := make([]bool, 0, len(txAcctMetas))
 	convertedAcctMetas := make([]*sealevel.AccountMeta, 0, len(txAcctMetas))
 	var loadedBytesAccumulator uint32
+	var loadedAcctCount uint64
+	var loadedAcctBytes uint64
 
 	for idx, acctMeta := range txAcctMetas {
 		var acct *accounts.Account
 		var isInstructionsSysvarAcct bool
+		var isSharedAcct bool
 
 		_, instrContainsAcctMeta := instructionAcctPubkeys[acctMeta.PublicKey]
 		if acctMeta.PublicKey == sealevel.SysvarInstructionsAddr {
 			acct = instrsAcct
 			isInstructionsSysvarAcct = true
 		} else if !slotCtx.Features.IsActive(features.DisableAccountLoaderSpecialCase) && slices.Contains(programIdIdxs, uint64(idx)) && !acctMeta.IsWritable && !instrContainsAcctMeta {
-			tmp, err := slotCtx.GetAccount(acctMeta.PublicKey)
+			tmp, err := slotCtx.GetAccountShared(acctMeta.PublicKey)
 			if err != nil {
 				return nil, nil, err
 			}
 			acct = &accounts.Account{Key: acctMeta.PublicKey, Owner: tmp.Owner, Executable: true, IsDummy: true}
 		} else {
-			acct, err = slotCtx.GetAccount(acctMeta.PublicKey)
+			acct, err = slotCtx.GetAccountShared(acctMeta.PublicKey)
 			if err != nil {
 				return nil, nil, err
 			}
+			isSharedAcct = true
 		}
 
 		if !isInstructionsSysvarAcct {
@@ -95,13 +113,21 @@ func loadAndValidateTxAccts(slotCtx *sealevel.SlotCtx, acctMetasPerInstr [][]sea
 			}
 		}
 
-		acctsForTx = append(acctsForTx, *acct)
+		acctsForTx = append(acctsForTx, acct)
+		acctsShared = append(acctsShared, isSharedAcct)
 		convertedAcctMeta := &sealevel.AccountMeta{Pubkey: acctMeta.PublicKey, IsSigner: acctMeta.IsSigner, IsWritable: acctMeta.IsWritable}
 		convertedAcctMetas = append(convertedAcctMetas, convertedAcctMeta)
+		if isSharedAcct {
+			loadedAcctCount++
+			loadedAcctBytes += uint64(len(acct.Data))
+		}
 	}
 
-	transactionAccts := sealevel.NewTransactionAccounts(acctsForTx)
+	transactionAccts := sealevel.NewTransactionAccountsFromRefs(acctsForTx, acctsShared)
 	transactionAccts.AcctMetas = convertedAcctMetas
+	transactionAccts.OnFirstWriteClone = recordTxAcctCowClone
+	TxAcctsLoaded.Add(loadedAcctCount)
+	TxAcctsLoadedBytes.Add(loadedAcctBytes)
 
 	removeAcctsExecutableFlagChecks := slotCtx.Features.IsActive(features.RemoveAccountsExecutableFlagChecks)
 	validatedLoaders := make(map[solana.PublicKey]struct{}, 4) // Usually ≤4 loaders
@@ -111,7 +137,7 @@ func loadAndValidateTxAccts(slotCtx *sealevel.SlotCtx, acctMetasPerInstr [][]sea
 			continue
 		}
 
-		programAcct, err := slotCtx.GetAccount(instr.ProgramId)
+		programAcct, err := slotCtx.GetAccountShared(instr.ProgramId)
 		if err != nil {
 			return nil, nil, TxErrProgramAccountNotFound
 		}
@@ -132,7 +158,7 @@ func loadAndValidateTxAccts(slotCtx *sealevel.SlotCtx, acctMetasPerInstr [][]sea
 		_, exists := validatedLoaders[owner]
 		if !exists {
 			var ownerAcct *accounts.Account
-			ownerAcct, err = slotCtx.GetAccount(owner)
+			ownerAcct, err = slotCtx.GetAccountShared(owner)
 			if err != nil {
 				ownerAcct, err = slotCtx.GetAccountFromAccountsDb(owner)
 				if err != nil {
@@ -216,7 +242,7 @@ func (accum *loadedAcctSizeAccumulatorSimd186) collectAcct(acct *accounts.Accoun
 		if err == nil && acctState.Type == sealevel.UpgradeableLoaderStateTypeProgram {
 			if !accum.wasAlreadyCounted(programDataAddr) {
 				// program data account not being found is not an error. Agave instead ignores it.
-				programDataAcct, err := accum.slotCtx.GetAccount(programDataAddr)
+				programDataAcct, err := accum.slotCtx.GetAccountShared(programDataAddr)
 				if err != nil {
 					programDataAcct, err = accum.slotCtx.GetAccountFromAccountsDb(programDataAddr)
 					if err != nil {
@@ -254,28 +280,27 @@ func loadAndValidateTxAcctsSimd186(slotCtx *sealevel.SlotCtx, acctMetasPerInstr
 		return nil, nil, err
 	}
 
-	// Memoize accounts loaded in Pass 1 to avoid re-cloning in Pass 2
+	// Memoize accounts loaded in Pass 1
 	// Use slice indexed by account position (same ordering as txAcctMetas)
 	acctCache := make([]*accounts.Account, len(acctKeys))
 
-	var clonedBytes uint64
 	for i, pubkey := range acctKeys {
-		acct, err := slotCtx.GetAccount(pubkey)
-		if err != nil {
-			panic("should be impossible - programming error")
+		var acct *accounts.Account
+		if pubkey == sealevel.SysvarInstructionsAddr {
+			acct = instrsAcct
+		} else {
+			acct, err = slotCtx.GetAccountShared(pubkey)
+			if err != nil {
+				panic("should be impossible - programming error")
+			}
 		}
 		acctCache[i] = acct // Cache by index for reuse in Pass 2
-		clonedBytes += uint64(len(acct.Data))
 		err = accumulator.collectAcct(acct)
 		if err != nil {
 			return nil, nil, err
 		}
 	}
 
-	// Track clone stats for profiling
-	TxAcctsCloned.Add(uint64(len(acctKeys)))
-	TxAcctsClonedBytes.Add(clonedBytes)
-
 	txAcctMetas, err := tx.AccountMetaList()
 	if err != nil {
 		return nil, nil, err
@@ -296,11 +321,15 @@ func loadAndValidateTxAcctsSimd186(slotCtx *sealevel.SlotCtx, acctMetasPerInstr
 		}
 	}
 
-	acctsForTx := make([]accounts.Account, 0, len(txAcctMetas))
+	acctsForTx := make([]*accounts.Account, 0, len(txAcctMetas))
+	acctsShared := make([]bool, 0, len(txAcctMetas))
 	convertedAcctMetas := make([]*sealevel.AccountMeta, 0, len(txAcctMetas))
+	var loadedAcctCount uint64
+	var loadedAcctBytes uint64
 
 	for idx, acctMeta := range txAcctMetas {
 		var acct *accounts.Account
+		var isSharedAcct bool
 		cached := acctCache[idx] // Reuse account from Pass 1
 
 		_, instrContainsAcctMeta := instructionAcctPubkeys[acctMeta.PublicKey]
@@ -312,15 +341,24 @@ func loadAndValidateTxAcctsSimd186(slotCtx *sealevel.SlotCtx, acctMetasPerInstr
 		} else {
 			// Normal case - use cached account directly
 			acct = cached
+			isSharedAcct = true
 		}
 
-		acctsForTx = append(acctsForTx, *acct)
+		acctsForTx = append(acctsForTx, acct)
+		acctsShared = append(acctsShared, isSharedAcct)
 		convertedAcctMeta := &sealevel.AccountMeta{Pubkey: acctMeta.PublicKey, IsSigner: acctMeta.IsSigner, IsWritable: acctMeta.IsWritable}
 		convertedAcctMetas = append(convertedAcctMetas, convertedAcctMeta)
+		if isSharedAcct {
+			loadedAcctCount++
+			loadedAcctBytes += uint64(len(acct.Data))
+		}
 	}
 
-	transactionAccts := sealevel.NewTransactionAccounts(acctsForTx)
+	transactionAccts := sealevel.NewTransactionAccountsFromRefs(acctsForTx, acctsShared)
 	transactionAccts.AcctMetas = convertedAcctMetas
+	transactionAccts.OnFirstWriteClone = recordTxAcctCowClone
+	TxAcctsLoaded.Add(loadedAcctCount)
+	TxAcctsLoadedBytes.Add(loadedAcctBytes)
 
 	removeAcctsExecutableFlagChecks := slotCtx.Features.IsActive(features.RemoveAccountsExecutableFlagChecks)
 
@@ -339,7 +377,7 @@ func loadAndValidateTxAcctsSimd186(slotCtx *sealevel.SlotCtx, acctMetasPerInstr
 		// Fallback if not in cache or out of bounds
 		if programAcct == nil {
 			var err error
-			programAcct, err = slotCtx.GetAccount(instr.ProgramId)
+			programAcct, err = slotCtx.GetAccountShared(instr.ProgramId)
 			if err != nil {
 				programAcct, err = slotCtx.GetAccountFromAccountsDb(instr.ProgramId)
 				if err != nil {

pkg/replay/block.go

@@ -2223,14 +2223,19 @@ func ReplayBlocks(
 				// Account clone stats for copy-on-write optimization profiling
 				cloneStats := GetAndResetCloneStats()
 				if cloneStats.TxCount > 0 {
-					modifyRatio := float64(cloneStats.AcctsTouched) / float64(cloneStats.AcctsCloned) * 100
-					avgAcctsPerTx := float64(cloneStats.AcctsCloned) / float64(cloneStats.TxCount)
+					var cloneRatio float64
+					if cloneStats.AcctsLoaded > 0 {
+						cloneRatio = float64(cloneStats.AcctsCloned) / float64(cloneStats.AcctsLoaded) * 100
+					}
+					avgLoadedPerTx := float64(cloneStats.AcctsLoaded) / float64(cloneStats.TxCount)
+					avgClonedPerTx := float64(cloneStats.AcctsCloned) / float64(cloneStats.TxCount)
 					avgTouchedPerTx := float64(cloneStats.AcctsTouched) / float64(cloneStats.TxCount)
+					loadedMB := float64(cloneStats.AcctsLoadedBytes) / 1024 / 1024
 					clonedMB := float64(cloneStats.AcctsClonedBytes) / 1024 / 1024
 					touchedMB := float64(cloneStats.AcctsTouchedBytes) / 1024 / 1024
-					mlog.Log.InfofPrecise("  clone stats: %.1f%% modified (%d/%d accts) | %.1fMB cloned, %.1fMB modified | avg/tx: %.1f cloned, %.1f modified",
-						modifyRatio, cloneStats.AcctsTouched, cloneStats.AcctsCloned,
-						clonedMB, touchedMB, avgAcctsPerTx, avgTouchedPerTx)
+					mlog.Log.InfofPrecise("  account COW: %.1f%% cloned on write (%d/%d accts) | %.1fMB loaded, %.1fMB cloned, %.1fMB modified | avg/tx: %.1f loaded, %.1f cloned, %.1f modified",
+						cloneRatio, cloneStats.AcctsCloned, cloneStats.AcctsLoaded,
+						loadedMB, clonedMB, touchedMB, avgLoadedPerTx, avgClonedPerTx, avgTouchedPerTx)
 				}
 
 				var mem runtime.MemStats

pkg/sealevel/borrowed_account.go

@@ -31,10 +31,11 @@ func (acct *BorrowedAccount) RentEpoch() uint64 {
 }
 
 func (acct *BorrowedAccount) Touch() error {
-	err := acct.TxCtx.Accounts.Touch(acct.IndexInTransaction)
+	touchedAcct, err := acct.TxCtx.Accounts.Touch(acct.IndexInTransaction)
 	if err != nil {
 		return err
 	}
+	acct.Account = touchedAcct
 	return nil
 }
 
@@ -269,7 +270,10 @@ func (acct *BorrowedAccount) SetDataLength(newLength uint64, f features.Features
 		return nil
 	}
 
-	acct.Touch()
+	err = acct.Touch()
+	if err != nil {
+		return err
+	}
 	acct.UpdateAccountsResizeDelta(newLength)
 	acct.Account.Resize(newLength, 0)
 

pkg/sealevel/execution_ctx.go

@@ -369,6 +369,11 @@ func (slotCtx *SlotCtx) GetAccount(pubkey solana.PublicKey) (*accounts.Account,
 	}
 }
 
+func (slotCtx *SlotCtx) GetAccountShared(pubkey solana.PublicKey) (*accounts.Account, error) {
+	pk := [32]byte(pubkey)
+	return slotCtx.Accounts.GetAccount(&pk)
+}
+
 func (slotCtx *SlotCtx) GetParentAccount(pubkey solana.PublicKey) (*accounts.Account, error) {
 	acct, err := slotCtx.ParentAccts.GetAccountWithoutLock(pubkey)
 	if err != nil {

pkg/sealevel/transaction_ctx.go

@@ -17,10 +17,12 @@ type TxReturnData struct {
 }
 
 type TransactionAccounts struct {
-	Accounts  []*accounts.Account
-	Locked    []bool
-	Touched   []bool
-	AcctMetas []*AccountMeta
+	Accounts          []*accounts.Account
+	Shared            []bool
+	Locked            []bool
+	Touched           []bool
+	AcctMetas         []*AccountMeta
+	OnFirstWriteClone func(*accounts.Account)
 }
 
 type TransactionCtx struct {
@@ -47,6 +49,7 @@ func NewTransactionAccounts(accts []accounts.Account) *TransactionAccounts {
 	transactionAccts := new(TransactionAccounts)
 
 	transactionAccts.Accounts = make([]*accounts.Account, 0, len(accts))
+	transactionAccts.Shared = make([]bool, len(accts))
 	for _, acct := range accts {
 		a := acct
 		transactionAccts.Accounts = append(transactionAccts.Accounts, &a)
@@ -58,6 +61,21 @@ func NewTransactionAccounts(accts []accounts.Account) *TransactionAccounts {
 	return transactionAccts
 }
 
+func NewTransactionAccountsFromRefs(accts []*accounts.Account, shared []bool) *TransactionAccounts {
+	if len(accts) != len(shared) {
+		panic("transaction accounts/shared flags length mismatch")
+	}
+
+	transactionAccts := new(TransactionAccounts)
+
+	transactionAccts.Accounts = append(make([]*accounts.Account, 0, len(accts)), accts...)
+	transactionAccts.Shared = append(make([]bool, 0, len(shared)), shared...)
+	transactionAccts.Locked = make([]bool, len(accts), len(accts))
+	transactionAccts.Touched = make([]bool, len(accts), len(accts))
+
+	return transactionAccts
+}
+
 func NewTransactionCtx(txAccts TransactionAccounts, instrStackCapacity uint64, instrTraceCapacity uint64) *TransactionCtx {
 	txCtx := new(TransactionCtx)
 
@@ -287,10 +305,20 @@ func (txAccounts *TransactionAccounts) Unlock(idx uint64) {
 	txAccounts.Locked[idx] = false
 }
 
-func (txAccounts *TransactionAccounts) Touch(idx uint64) error {
+func (txAccounts *TransactionAccounts) Touch(idx uint64) (*accounts.Account, error) {
 	if len(txAccounts.Touched) == 0 || idx > uint64(len(txAccounts.Touched)-1) {
-		return InstrErrNotEnoughAccountKeys
+		return nil, InstrErrNotEnoughAccountKeys
+	}
+
+	if txAccounts.Shared[idx] {
+		clonedAcct := txAccounts.Accounts[idx].Clone()
+		txAccounts.Accounts[idx] = clonedAcct
+		txAccounts.Shared[idx] = false
+		if txAccounts.OnFirstWriteClone != nil {
+			txAccounts.OnFirstWriteClone(clonedAcct)
+		}
 	}
+
 	txAccounts.Touched[idx] = true
-	return nil
+	return txAccounts.Accounts[idx], nil
 }