Ordered account insertion into merkle trie (I)

ledger/accountdb.go

@@ -1131,3 +1131,236 @@ func (iterator *encodedAccountsBatchIter) Close() {
 		iterator.rows = nil
 	}
 }
+
+// orderedAccountsIterStep is used by orderedAccountsIter to define the current step
+//msgp:ignore orderedAccountsIterStep
+type orderedAccountsIterStep int
+
+const (
+	// startup step
+	oaiStepStartup = orderedAccountsIterStep(0)
+	// delete old ordering table if we have any leftover from previous invocation
+	oaiStepDeleteOldOrderingTable = orderedAccountsIterStep(0)
+	// create new ordering table
+	oaiStepCreateOrderingTable = orderedAccountsIterStep(1)
+	// query the existing accounts
+	oaiStepQueryAccounts = orderedAccountsIterStep(2)
+	// iterate over the existing accounts and insert their hash & address into the staging ordering table
+	oaiStepInsertAccountData = orderedAccountsIterStep(3)
+	// create an index on the ordering table so that we can efficiently scan it.
+	oaiStepCreateOrderingAccountIndex = orderedAccountsIterStep(4)
+	// query the ordering table
+	oaiStepSelectFromOrderedTable = orderedAccountsIterStep(5)
+	// iterate over the ordering table
+	oaiStepIterateOverOrderedTable = orderedAccountsIterStep(6)
+	// cleanup and delete ordering table
+	oaiStepShutdown = orderedAccountsIterStep(7)
+	// do nothing as we're done.
+	oaiStepDone = orderedAccountsIterStep(8)
+)
+
+// orderedAccountsIter allows us to iterate over the accounts addresses in the order of the account hashes.
+type orderedAccountsIter struct {
+	step             orderedAccountsIterStep
+	rows             *sql.Rows
+	tx               *sql.Tx
+	accountCount     int
+	fetchAccountData bool
+	insertStmt       *sql.Stmt
+}
+
+// makeOrderedAccountsIter creates an ordered account iterator. Note that due to implementation reasons,
+// only a single iterator can be active at a time.
+func makeOrderedAccountsIter(tx *sql.Tx, accountCount int, fetchAccountData bool) *orderedAccountsIter {
+	return &orderedAccountsIter{
+		tx:               tx,
+		accountCount:     accountCount,
+		fetchAccountData: fetchAccountData,
+		step:             oaiStepStartup,
+	}
+}
+
+// accountAddressHashData is used by Next to return a single account address, hash and account data.
+type accountAddressHashData struct {
+	address            basics.Address
+	digest             []byte
+	encodedAccountData []byte
+}
+
+// Next returns an array containing the account address, hash and potentially data
+// the Next function works in multiple processing stages, where it first processs the current accounts and order them
+// followed by returning the ordered accounts. In the first phase, it would return empty accountAddressHashData array
+// and sets the processedRecords to the number of accounts that were processed. On the second phase, the acct
+// would contain valid data ( and optionally the account data as well, if was asked in makeOrderedAccountsIter) and
+// the processedRecords would be zero. If err is sql.ErrNoRows it means that the iterator have completed it's work and no further
+// accounts exists. Otherwise, the caller is expected to keep calling "Next" to retrieve the next set of accounts
+// ( or let the Next function make some progress toward that goal )
+func (iterator *orderedAccountsIter) Next(ctx context.Context) (acct []accountAddressHashData, processedRecords int, err error) {
+	if iterator.step == oaiStepDeleteOldOrderingTable {
+		// although we're going to delete this table anyway when completing the iterator execution, we'll try to
+		// clean up any intermediate table.
+		_, err = iterator.tx.ExecContext(ctx, "DROP TABLE IF EXISTS accountsiteratorhashes")
+		if err != nil {
+			return
+		}
+		iterator.step = oaiStepCreateOrderingTable
+		return
+	}
+	if iterator.step == oaiStepCreateOrderingTable {
+		// create the temporary table
+		_, err = iterator.tx.ExecContext(ctx, "CREATE TABLE accountsiteratorhashes(address blob, hash blob)")
+		if err != nil {
+			return
+		}
+		iterator.step = oaiStepQueryAccounts
+		return
+	}
+	if iterator.step == oaiStepQueryAccounts {
+		// iterate over the existing accounts
+		iterator.rows, err = iterator.tx.QueryContext(ctx, "SELECT address, data FROM accountbase")
+		if err != nil {
+			return
+		}
+		// prepare the insert statement into the temporary table
+		iterator.insertStmt, err = iterator.tx.PrepareContext(ctx, "INSERT INTO accountsiteratorhashes(address, hash) VALUES(?, ?)")
+		if err != nil {
+			return
+		}
+		iterator.step = oaiStepInsertAccountData
+		return
+	}
+	if iterator.step == oaiStepInsertAccountData {
+		var addr basics.Address
+		count := 0
+		for iterator.rows.Next() {
+			var addrbuf []byte
+			var buf []byte
+			err = iterator.rows.Scan(&addrbuf, &buf)
+			if err != nil {
+				iterator.Close(ctx)
+				return
+			}
+
+			if len(addrbuf) != len(addr) {
+				err = fmt.Errorf("Account DB address length mismatch: %d != %d", len(addrbuf), len(addr))
+				iterator.Close(ctx)
+				return
+			}
+
+			copy(addr[:], addrbuf)
+
+			var accountData basics.AccountData
+			err = protocol.Decode(buf, &accountData)
+			if err != nil {
+				iterator.Close(ctx)
+				return
+			}
+			hash := accountHashBuilder(addr, accountData, buf)
+			_, err = iterator.insertStmt.ExecContext(ctx, addrbuf, hash)
+			if err != nil {
+				iterator.Close(ctx)
+				return
+			}
+
+			count++
+			if count == iterator.accountCount {
+				// we're done with this iteration.
+				processedRecords = count
+				return
+			}
+		}
+		processedRecords = count
+		iterator.rows.Close()
+		iterator.rows = nil
+		iterator.insertStmt.Close()
+		iterator.insertStmt = nil
+		iterator.step = oaiStepCreateOrderingAccountIndex
+		return
+	}
+	if iterator.step == oaiStepCreateOrderingAccountIndex {
+		// create an index. It shown that even when we're making a single select statement in step 5, it would be better to have this index vs. not having it at all.
+		// note that this index is using the rowid of the accountsiteratorhashes table.
+		_, err = iterator.tx.ExecContext(ctx, "CREATE INDEX accountsiteratorhashesidx ON accountsiteratorhashes(hash)")
+		if err != nil {
+			iterator.Close(ctx)
+			return
+		}
+		iterator.step = oaiStepSelectFromOrderedTable
+		return
+	}
+	if iterator.step == oaiStepSelectFromOrderedTable {
+		// select the data from the ordered table
+		if iterator.fetchAccountData {
+			iterator.rows, err = iterator.tx.QueryContext(ctx, "SELECT accountsiteratorhashes.address, accountsiteratorhashes.hash, accountbase.data FROM accountsiteratorhashes JOIN accountbase ON accountbase.address=accountsiteratorhashes.address ORDER BY accountsiteratorhashes.hash")
+		} else {
+			iterator.rows, err = iterator.tx.QueryContext(ctx, "SELECT address, hash FROM accountsiteratorhashes ORDER BY hash")
+		}
+
+		if err != nil {
+			iterator.Close(ctx)
+			return
+		}
+		iterator.step = oaiStepIterateOverOrderedTable
+		return
+	}
+
+	if iterator.step == oaiStepIterateOverOrderedTable {
+		acct = make([]accountAddressHashData, 0, iterator.accountCount)
+		var addr basics.Address
+		for iterator.rows.Next() {
+			var addrbuf []byte
+			var acctdata []byte
+			var hash []byte
+			if iterator.fetchAccountData {
+				err = iterator.rows.Scan(&addrbuf, &hash, &acctdata)
+			} else {
+				err = iterator.rows.Scan(&addrbuf, &hash)
+			}
+			if err != nil {
+				iterator.Close(ctx)
+				return
+			}
+
+			if len(addrbuf) != len(addr) {
+				err = fmt.Errorf("Account DB address length mismatch: %d != %d", len(addrbuf), len(addr))
+				iterator.Close(ctx)
+				return
+			}
+
+			copy(addr[:], addrbuf)
+
+			acct = append(acct, accountAddressHashData{address: addr, digest: hash, encodedAccountData: acctdata})
+			if len(acct) == iterator.accountCount {
+				// we're done with this iteration.
+				return
+			}
+		}
+		iterator.step = oaiStepShutdown
+		iterator.rows.Close()
+		iterator.rows = nil
+		return
+	}
+	if iterator.step == oaiStepShutdown {
+		err = iterator.Close(ctx)
+		if err != nil {
+			return
+		}
+		iterator.step = oaiStepDone
+		// fallthrough
+	}
+	return nil, 0, sql.ErrNoRows
+}
+
+// Close shuts down the orderedAccountsBuilderIter, releasing database resources.
+func (iterator *orderedAccountsIter) Close(ctx context.Context) (err error) {
+	if iterator.rows != nil {
+		iterator.rows.Close()
+		iterator.rows = nil
+	}
+	if iterator.insertStmt != nil {
+		iterator.insertStmt.Close()
+		iterator.insertStmt = nil
+	}
+	_, err = iterator.tx.ExecContext(ctx, "DROP TABLE IF EXISTS accountsiteratorhashes")
+	return
+}

ledger/acctupdates.go

@@ -1119,56 +1119,58 @@ func (au *accountUpdates) accountsInitialize(ctx context.Context, tx *sql.Tx) (b
 
 	if rootHash.IsZero() {
 		au.log.Infof("accountsInitialize rebuilding merkle trie for round %d", rnd)
-		var accountsIterator encodedAccountsBatchIter
-		defer accountsIterator.Close()
+		accountBuilderIt := makeOrderedAccountsIter(tx, trieRebuildAccountChunkSize, false)
+		defer accountBuilderIt.Close(ctx)
 		startTrieBuildTime := time.Now()
 		accountsCount := 0
 		lastRebuildTime := startTrieBuildTime
 		pendingAccounts := 0
+		totalOrderedAccounts := 0
 		for {
-			bal, err := accountsIterator.Next(ctx, tx, trieRebuildAccountChunkSize)
-			if err != nil {
-				return rnd, err
-			}
-			if len(bal) == 0 {
+			accts, processedRows, err := accountBuilderIt.Next(ctx)
+			if err == sql.ErrNoRows {
+				// the account builder would return sql.ErrNoRows when no more data is available.
 				break
-			}
-			accountsCount += len(bal)
-			pendingAccounts += len(bal)
-			for _, balance := range bal {
-				var accountData basics.AccountData
-				err = protocol.Decode(balance.AccountData, &accountData)
-				if err != nil {
-					return rnd, err
-				}
-				hash := accountHashBuilder(balance.Address, accountData, balance.AccountData)
-				added, err := trie.Add(hash)
-				if err != nil {
-					return rnd, fmt.Errorf("accountsInitialize was unable to add changes to trie: %v", err)
-				}
-				if !added {
-					au.log.Warnf("accountsInitialize attempted to add duplicate hash '%s' to merkle trie for account %v", hex.EncodeToString(hash), balance.Address)
-				}
+			} else if err != nil {
+				return rnd, err
 			}
 
-			if pendingAccounts >= trieRebuildCommitFrequency {
-				// this trie Evict will commit using the current transaction.
-				// if anything goes wrong, it will still get rolled back.
-				_, err = trie.Evict(true)
-				if err != nil {
-					return 0, fmt.Errorf("accountsInitialize was unable to commit changes to trie: %v", err)
+			if len(accts) > 0 {
+				accountsCount += len(accts)
+				pendingAccounts += len(accts)
+				for _, acct := range accts {
+					added, err := trie.Add(acct.digest)
+					if err != nil {
+						return rnd, fmt.Errorf("accountsInitialize was unable to add changes to trie: %v", err)
+					}
+					if !added {
+						au.log.Warnf("accountsInitialize attempted to add duplicate hash '%s' to merkle trie for account %v", hex.EncodeToString(acct.digest), acct.address)
+					}
 				}
-				pendingAccounts = 0
-			}
 
-			if len(bal) < trieRebuildAccountChunkSize {
-				break
-			}
+				if pendingAccounts >= trieRebuildCommitFrequency {
+					// this trie Evict will commit using the current transaction.
+					// if anything goes wrong, it will still get rolled back.
+					_, err = trie.Evict(true)
+					if err != nil {
+						return 0, fmt.Errorf("accountsInitialize was unable to commit changes to trie: %v", err)
+					}
+					pendingAccounts = 0
+				}
 
-			if time.Now().Sub(lastRebuildTime) > 5*time.Second {
-				// let the user know that the trie is still being rebuilt.
-				au.log.Infof("accountsInitialize still building the trie, and processed so far %d accounts", accountsCount)
-				lastRebuildTime = time.Now()
+				if time.Now().Sub(lastRebuildTime) > 5*time.Second {
+					// let the user know that the trie is still being rebuilt.
+					au.log.Infof("accountsInitialize still building the trie, and processed so far %d accounts", accountsCount)
+					lastRebuildTime = time.Now()
+				}
+			} else if processedRows > 0 {
+				totalOrderedAccounts += processedRows
+				// if it's not ordered, we can ignore it for now; we'll just increase the counters and emit logs periodically.
+				if time.Now().Sub(lastRebuildTime) > 5*time.Second {
+					// let the user know that the trie is still being rebuilt.
+					au.log.Infof("accountsInitialize still building the trie, and hashed so far %d accounts", totalOrderedAccounts)
+					lastRebuildTime = time.Now()
+				}
 			}
 		}