From 3a0abc4a74d1a7acc2d1a779a07eda29f8ff5896 Mon Sep 17 00:00:00 2001
From: John Jannotti <jannotti@gmail.com>
Date: Wed, 27 Nov 2024 13:30:31 -0500
Subject: [PATCH] Prefetching for heartbeat transactions

---
 ledger/eval/prefetcher/prefetcher.go          |  4 +-
 .../prefetcher/prefetcher_alignment_test.go   | 52 +++++++++++++++++++
 2 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/ledger/eval/prefetcher/prefetcher.go b/ledger/eval/prefetcher/prefetcher.go
index 487d370524..765b6ea9c2 100644
--- a/ledger/eval/prefetcher/prefetcher.go
+++ b/ledger/eval/prefetcher/prefetcher.go
@@ -343,7 +343,9 @@ func (p *accountPrefetcher) prefetch(ctx context.Context) {
 				// since they might be non-used arbitrary values
 
 			case protocol.StateProofTx:
-			case protocol.KeyRegistrationTx:
+			case protocol.KeyRegistrationTx: // No extra accounts besides the sender
+			case protocol.HeartbeatTx:
+				loadAccountsAddAccountTask(&stxn.Txn.HbAddress, task, accountTasks, queue)
 			}
 
 			// If you add new addresses here, also add them in getTxnAddresses().
diff --git a/ledger/eval/prefetcher/prefetcher_alignment_test.go b/ledger/eval/prefetcher/prefetcher_alignment_test.go
index 5f61f4938f..6cd6dc9247 100644
--- a/ledger/eval/prefetcher/prefetcher_alignment_test.go
+++ b/ledger/eval/prefetcher/prefetcher_alignment_test.go
@@ -30,6 +30,7 @@ import (
 	"github.com/algorand/go-algorand/crypto/stateproof"
 	"github.com/algorand/go-algorand/data/basics"
 	"github.com/algorand/go-algorand/data/bookkeeping"
+	"github.com/algorand/go-algorand/data/committee"
 	"github.com/algorand/go-algorand/data/stateproofmsg"
 	"github.com/algorand/go-algorand/data/transactions"
 	"github.com/algorand/go-algorand/ledger/eval"
@@ -1422,3 +1423,54 @@ func TestEvaluatorPrefetcherAlignmentStateProof(t *testing.T) {
 	prefetched.pretend(rewardsPool())
 	require.Equal(t, requested, prefetched)
 }
+
+func TestEvaluatorPrefetcherAlignmentHeartbeat(t *testing.T) {
+	partitiontest.PartitionTest(t)
+
+	// We need valid part keys to evaluate the Heartbeat.
+	const kd = 10
+	firstID := basics.OneTimeIDForRound(0, kd)
+	otss := crypto.GenerateOneTimeSignatureSecrets(firstID.Batch, 5)
+
+	l := &prefetcherAlignmentTestLedger{
+		balances: map[basics.Address]ledgercore.AccountData{
+			rewardsPool(): {
+				AccountBaseData: ledgercore.AccountBaseData{
+					MicroAlgos: basics.MicroAlgos{Raw: 1234567890},
+				},
+			},
+			makeAddress(1): {
+				AccountBaseData: ledgercore.AccountBaseData{
+					MicroAlgos: basics.MicroAlgos{Raw: 1000001},
+				},
+			},
+			makeAddress(2): {
+				AccountBaseData: ledgercore.AccountBaseData{
+					MicroAlgos: basics.MicroAlgos{Raw: 100_000},
+				},
+				VotingData: basics.VotingData{
+					VoteID: otss.OneTimeSignatureVerifier,
+				},
+			},
+		},
+	}
+
+	txn := transactions.Transaction{
+		Type: protocol.HeartbeatTx,
+		Header: transactions.Header{
+			Sender:      makeAddress(1),
+			GenesisHash: genesisHash(),
+			Fee:         basics.Algos(1), // Heartbeat txn is unusual in that it checks fees a bit.
+		},
+		HeartbeatTxnFields: transactions.HeartbeatTxnFields{
+			HbAddress: makeAddress(2),
+			HbProof:   otss.Sign(firstID, committee.Seed(genesisHash())).ToHeartbeatProof(),
+			HbSeed:    committee.Seed(genesisHash()),
+		},
+	}
+
+	requested, prefetched := run(t, l, txn)
+
+	prefetched.pretend(rewardsPool())
+	require.Equal(t, requested, prefetched)
+}