blockprod: resolve address lookup tables for v0 transactions

pkg/blockprod/bank.go

@@ -152,6 +152,7 @@ const (
 	ForgeDroppedExecution
 	ForgeDroppedAlreadyProcessed
 	ForgeDroppedAccountLocks
+	ForgeDroppedLookup
 )
 
 func (r ForgeResult) String() string {
@@ -172,6 +173,8 @@ func (r ForgeResult) String() string {
 		return "dropped_already_processed"
 	case ForgeDroppedAccountLocks:
 		return "dropped_account_locks"
+	case ForgeDroppedLookup:
+		return "dropped_lookup"
 	default:
 		return "unknown"
 	}
@@ -246,8 +249,29 @@ func (b *WorkingBank) forgeTransaction(tx *solana.Transaction, wireSize int, pre
 		ancestorAlreadyProcessed, ancestorStatusErr = b.ancestorStatuses.ContainsMessage(tx.Message.RecentBlockhash, messageHash)
 	}
 
+	// For a v0 transaction with address-table lookups, execTx is a resolved copy,
+	// while tx stays unresolved for entry encoding.
+	execTx := tx
 	var cost costmodel.TransactionCost
 	if !ancestorAlreadyProcessed && ancestorStatusErr == nil {
+		if replay.HasAddressTableLookups(tx) {
+			// Tables are read from the parent state, which this slot does not
+			// change, so resolve outside the bank lock.
+			resolved, txErr := replay.ResolveAddressTableLookups(b.slotCtx, tx)
+			if txErr != nil {
+				b.RebateSchedule(wireSize)
+				return ForgeDroppedLookup, costmodel.ExceedNone
+			}
+			// A loaded address can repeat a static key or another loaded
+			// address, or take the account count past the lock limit.
+			if txErr := replay.ValidateAccountLocks(resolved, lockFeats); txErr != nil {
+				b.RebateSchedule(wireSize)
+				return ForgeDroppedAccountLocks, costmodel.ExceedNone
+			}
+			execTx = resolved
+			// Static preparation, if any, was done on the unresolved message.
+			prepared = nil
+		}
 		feats := b.slotCtx.Features
 		if feats == nil {
 			f := features.NewFeaturesDefault()
@@ -256,7 +280,7 @@ func (b *WorkingBank) forgeTransaction(tx *solana.Transaction, wireSize int, pre
 		if prepared != nil {
 			cost = prepared.Cost()
 		} else {
-			cost, err = costmodel.EstimateTransactionCost(tx, feats)
+			cost, err = costmodel.EstimateTransactionCost(execTx, feats)
 		}
 		if err != nil {
 			b.RebateSchedule(wireSize)
@@ -294,19 +318,19 @@ func (b *WorkingBank) forgeTransaction(tx *solana.Transaction, wireSize int, pre
 	if reason := b.reserveEntryBytesLocked(wireSize); reason != costmodel.ExceedNone {
 		return ForgeDroppedCost, reason
 	}
-	if err := b.preparer.PayerCanFund(b.slotCtx, tx, prepared); err != nil {
+	if err := b.preparer.PayerCanFund(b.slotCtx, execTx, prepared); err != nil {
 		return ForgeDroppedExecution, costmodel.ExceedNone
 	}
 
 	output := b.preparer.LoadAndExecute(replay.LoadAndExecuteTransactionInput{
 		SlotCtx:           b.slotCtx,
-		Transaction:       tx,
+		Transaction:       execTx,
 		LeanResult:        true,
 		SkipTimingMetrics: true,
 		Arena:             b.borrowedAccounts,
 	}, prepared)
 	if output.ProcessingResult.TransactionError != nil {
-		feeInfo, err := replay.ApplyFailedTransaction(b.slotCtx, tx, output)
+		feeInfo, err := replay.ApplyFailedTransaction(b.slotCtx, execTx, output)
 		if err != nil {
 			return ForgeDroppedExecution, costmodel.ExceedNone
 		}

pkg/blockprod/bank_lookup_test.go

@@ -0,0 +1,139 @@
+package blockprod
+
+import (
+	"testing"
+
+	"github.com/Overclock-Validator/mithril/pkg/costmodel"
+	"github.com/Overclock-Validator/mithril/pkg/sealevel"
+	"github.com/Overclock-Validator/mithril/pkg/tpu/txfixture"
+	"github.com/gagliardetto/solana-go"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+var lookupTestTableKey = solana.PublicKey{0x7A, 0xB1, 0xE0}
+
+// withLookupSlotHashes installs SlotHashes for the test env's slot, which
+// lookup resolution needs to classify deactivating tables.
+func withLookupSlotHashes(t *testing.T, env *TestEnv) {
+	t.Helper()
+	prev := sealevel.SysvarCache.SlotHashes.Sysvar
+	slot := env.SlotCtx.Slot
+	slotHashes := sealevel.SysvarSlotHashes{{Slot: slot - 1}, {Slot: slot - 2}}
+	sealevel.SysvarCache.SlotHashes.Sysvar = &slotHashes
+	t.Cleanup(func() { sealevel.SysvarCache.SlotHashes.Sysvar = prev })
+}
+
+// addLookupTable puts table in the test bank's parent state, which is where
+// lookups are resolved from.
+func addLookupTable(t *testing.T, env *TestEnv, table LookupTable) {
+	t.Helper()
+	acct, err := table.Account()
+	require.NoError(t, err)
+	parent, ok := env.SlotCtx.UnrootedRead.(ParentAccounts)
+	if !ok {
+		parent = ParentAccounts{}
+		env.SlotCtx.UnrootedRead = parent
+	}
+	parent[table.Key] = acct
+}
+
+func lookupTransfer(t *testing.T, lamports uint64) (*solana.Transaction, []byte) {
+	t.Helper()
+	wire, err := SignedV0LookupTransferWire(lamports, lookupTestTableKey, 0)
+	require.NoError(t, err)
+	tx, err := solana.TransactionFromBytes(wire)
+	require.NoError(t, err)
+	return tx, wire
+}
+
+func lamportsOf(t *testing.T, env *TestEnv, key solana.PublicKey) uint64 {
+	t.Helper()
+	acct, err := env.SlotCtx.GetAccount(key)
+	require.NoError(t, err)
+	return acct.Lamports
+}
+
+func TestWorkingBankForgesV0TransferThroughLookupTable(t *testing.T) {
+	env := NewTestEnv(TestEnvConfig{})
+	defer env.Close()
+	withLookupSlotHashes(t, env)
+	addLookupTable(t, env, ActiveLookupTable(lookupTestTableKey, txfixture.DestPubkey()))
+
+	const amount = 1_234
+	tx, wire := lookupTransfer(t, amount)
+	payerBefore := lamportsOf(t, env, txfixture.PayerPubkey())
+	destBefore := lamportsOf(t, env, txfixture.DestPubkey())
+
+	result, reason := env.Bank.ForgeTransaction(tx, len(wire))
+	require.Equal(t, ForgeAccepted, result)
+	require.Equal(t, costmodel.ExceedNone, reason)
+
+	assert.Equal(t, destBefore+amount, lamportsOf(t, env, txfixture.DestPubkey()))
+	assert.Equal(t, payerBefore-amount-5_000, lamportsOf(t, env, txfixture.PayerPubkey()))
+	assert.Equal(t, uint64(5_000), env.Bank.TxFeeAccumulator().TotalFees)
+
+	// The entry carries the original, unresolved transaction and its exact wire.
+	forged := env.Bank.ForgedTransactions()
+	require.Len(t, forged, 1)
+	assert.Same(t, tx, forged[0])
+	assert.False(t, forged[0].Message.IsResolved())
+	forgedWire, err := forged[0].MarshalBinary()
+	require.NoError(t, err)
+	assert.Equal(t, wire, forgedWire)
+
+	// The looked-up destination is write-locked in block cost accounting.
+	assert.NotZero(t, env.Bank.CostTracker().WritableAccountCost(txfixture.DestPubkey()))
+}
+
+func TestWorkingBankDropsUnresolvableLookups(t *testing.T) {
+	deactivated := ActiveLookupTable(lookupTestTableKey, txfixture.DestPubkey())
+	deactivated.DeactivationSlot = 30 // no longer in SlotHashes
+	appendedThisSlot := ActiveLookupTable(lookupTestTableKey, txfixture.DestPubkey())
+	appendedThisSlot.LastExtendedSlot = 42 // the test env's slot
+	appendedThisSlot.LastExtendedStartIndex = 0
+
+	cases := []struct {
+		name  string
+		table *LookupTable
+	}{
+		{"missing table", nil},
+		{"deactivated table", &deactivated},
+		{"address appended this slot", &appendedThisSlot},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			env := NewTestEnv(TestEnvConfig{})
+			defer env.Close()
+			withLookupSlotHashes(t, env)
+			if tc.table != nil {
+				addLookupTable(t, env, *tc.table)
+			}
+			tx, wire := lookupTransfer(t, 1)
+			assertForgeLookupDropped(t, env, tx, wire, ForgeDroppedLookup)
+		})
+	}
+}
+
+// The account lock rules apply to the resolved keys too: a table that loads
+// the transaction's own payer repeats a static key.
+func TestWorkingBankDropsLookupThatRepeatsAStaticKey(t *testing.T) {
+	env := NewTestEnv(TestEnvConfig{})
+	defer env.Close()
+	withLookupSlotHashes(t, env)
+	addLookupTable(t, env, ActiveLookupTable(lookupTestTableKey, txfixture.PayerPubkey()))
+	tx, wire := lookupTransfer(t, 1)
+	assertForgeLookupDropped(t, env, tx, wire, ForgeDroppedAccountLocks)
+}
+
+func assertForgeLookupDropped(t *testing.T, env *TestEnv, tx *solana.Transaction, wire []byte, want ForgeResult) {
+	t.Helper()
+	payerBefore := lamportsOf(t, env, txfixture.PayerPubkey())
+	result, reason := env.Bank.ForgeTransaction(tx, len(wire))
+	assert.Equal(t, want, result)
+	assert.Equal(t, costmodel.ExceedNone, reason)
+	assert.Equal(t, payerBefore, lamportsOf(t, env, txfixture.PayerPubkey()))
+	assert.Empty(t, env.Bank.ForgedTransactions())
+	assert.Zero(t, env.Bank.TxFeeAccumulator().TotalFees)
+	assert.False(t, tx.Message.IsResolved(), "the scheduler's transaction must stay unresolved")
+}

pkg/blockprod/scheduler/scheduler.go

@@ -32,6 +32,7 @@ type Stats struct {
 	DroppedExpired          uint64
 	DroppedAlreadyProcessed uint64
 	DroppedAccountLocks     uint64
+	DroppedLookup           uint64
 	Accepted                uint64
 	DroppedNoBank           uint64
 	DroppedCost             uint64
@@ -365,6 +366,8 @@ func (s *Scheduler) recordForge(result blockprod.ForgeResult, reason costmodel.E
 		s.stats.DroppedAlreadyProcessed++
 	case blockprod.ForgeDroppedAccountLocks:
 		s.stats.DroppedAccountLocks++
+	case blockprod.ForgeDroppedLookup:
+		s.stats.DroppedLookup++
 	default:
 		s.stats.DroppedNoBank++
 	}

pkg/blockprod/scheduler/scheduler_lookup_test.go

@@ -0,0 +1,73 @@
+package scheduler
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/Overclock-Validator/mithril/pkg/blockprod"
+	"github.com/Overclock-Validator/mithril/pkg/sealevel"
+	"github.com/Overclock-Validator/mithril/pkg/tpu/packet"
+	"github.com/Overclock-Validator/mithril/pkg/tpu/txfixture"
+	"github.com/gagliardetto/solana-go"
+	"github.com/stretchr/testify/require"
+)
+
+var schedulerLookupTableKey = solana.PublicKey{0x7A, 0xB1, 0xE1}
+
+// lookupEnv is a test bank at slot whose lookup table had its only address
+// appended in slot 42.
+func lookupEnv(t *testing.T, slot uint64) *blockprod.TestEnv {
+	t.Helper()
+	env := blockprod.NewTestEnv(blockprod.TestEnvConfig{})
+	t.Cleanup(env.Close)
+	env.SlotCtx.Slot = slot
+	table := blockprod.ActiveLookupTable(schedulerLookupTableKey, txfixture.DestPubkey())
+	table.LastExtendedSlot = 42
+	acct, err := table.Account()
+	require.NoError(t, err)
+	env.SlotCtx.UnrootedRead = blockprod.ParentAccounts{table.Key: acct}
+	return env
+}
+
+func waitFor(t *testing.T, what string, cond func() bool) {
+	t.Helper()
+	deadline := time.Now().Add(2 * time.Second)
+	for !cond() {
+		if time.Now().After(deadline) {
+			t.Fatalf("timed out waiting for %s", what)
+		}
+		time.Sleep(2 * time.Millisecond)
+	}
+}
+
+// A transaction whose lookup does not resolve is dropped, not retried, even
+// though it would resolve in a later slot.
+func TestSchedulerDropsFailedLookupWithoutRetry(t *testing.T) {
+	prev := sealevel.SysvarCache.SlotHashes.Sysvar
+	slotHashes := sealevel.SysvarSlotHashes{{Slot: 42}, {Slot: 41}}
+	sealevel.SysvarCache.SlotHashes.Sysvar = &slotHashes
+	t.Cleanup(func() { sealevel.SysvarCache.SlotHashes.Sysvar = prev })
+
+	controller := blockprod.NewController()
+	sched := New(controller)
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+	sched.Start(ctx)
+	t.Cleanup(sched.Stop)
+	wire, err := blockprod.SignedV0LookupTransferWire(1, schedulerLookupTableKey, 0)
+	require.NoError(t, err)
+	sched.Receive(packet.Owned(wire))
+	waitFor(t, "buffering", func() bool { return sched.Buffered() == 1 })
+
+	controller.SetWorkingBank(lookupEnv(t, 42).Bank)
+	waitFor(t, "lookup failure", func() bool { return sched.Stats().DroppedLookup == 1 })
+
+	next := lookupEnv(t, 43)
+	controller.SetWorkingBank(next.Bank)
+	time.Sleep(50 * time.Millisecond)
+	require.Zero(t, sched.Stats().Accepted)
+	require.Empty(t, next.Bank.ForgedTransactions())
+	require.Zero(t, sched.Buffered())
+	require.Equal(t, uint64(1), sched.Stats().DroppedLookup)
+}