Make missing txs check atomic

Author: codchen

internal/consensus/replay_stubs.go

@@ -24,6 +24,10 @@ func (m emptyMempool) GetTxsForKeys(txKeys []types.TxKey) types.Txs {
 	return types.Txs{}
 }
 
+func (m emptyMempool) SafeGetTxsForKeys(txKeys []types.TxKey) (types.Txs, []types.TxKey) {
+	return types.Txs{}, []types.TxKey{}
+}
+
 var _ mempool.Mempool = emptyMempool{}
 
 func (emptyMempool) TxStore() *mempool.TxStore { return nil }

internal/consensus/state.go

@@ -2502,13 +2502,12 @@ func (cs *State) tryCreateProposalBlock(ctx context.Context, height int64, round
 // Build a proposal block from mempool txs. If cs.config.GossipTransactionKeyOnly=true
 // proposals only contain txKeys so we rebuild the block using mempool txs
 func (cs *State) buildProposalBlock(height int64, header types.Header, lastCommit *types.Commit, evidence []types.Evidence, proposerAddress types.Address, txKeys []types.TxKey) *types.Block {
-	missingTxs := cs.blockExec.GetMissingTxs(txKeys)
+	txs, missingTxs := cs.blockExec.SafeGetTxsByKeys(txKeys)
 	if len(missingTxs) > 0 {
 		cs.metrics.ProposalMissingTxs.Set(float64(len(missingTxs)))
 		cs.logger.Debug("Missing txs when trying to build block", "missing_txs", cs.blockExec.GetMissingTxs(txKeys))
 		return nil
 	}
-	txs := cs.blockExec.GetTxsForKeys(txKeys)
 	block := cs.state.MakeBlock(height, cs.blockExec.GetTxsForKeys(txKeys), lastCommit, evidence, proposerAddress)
 	block.Version = header.Version
 	block.Data.Txs = txs

internal/mempool/mempool.go

@@ -405,6 +405,23 @@ func (txmp *TxMempool) GetTxsForKeys(txKeys []types.TxKey) types.Txs {
 	return txs
 }
 
+func (txmp *TxMempool) SafeGetTxsForKeys(txKeys []types.TxKey) (types.Txs, []types.TxKey) {
+	txmp.mtx.RLock()
+	defer txmp.mtx.RUnlock()
+
+	txs := make([]types.Tx, 0, len(txKeys))
+	missing := []types.TxKey{}
+	for _, txKey := range txKeys {
+		wtx := txmp.txStore.GetTxByHash(txKey)
+		if wtx == nil {
+			missing = append(missing, txKey)
+			continue
+		}
+		txs = append(txs, wtx.tx)
+	}
+	return txs, missing
+}
+
 // Flush empties the mempool. It acquires a read-lock, fetches all the
 // transactions currently in the transaction store and removes each transaction
 // from the store and all indexes and finally resets the cache.

internal/mempool/mocks/mempool.go

@@ -42,6 +42,10 @@ type Mempool interface {
 
 	GetTxsForKeys(txKeys []types.TxKey) types.Txs
 
+	// Similar to GetTxsForKeys except that it would return a list
+	// indicating missing keys.
+	SafeGetTxsForKeys(txKeys []types.TxKey) (types.Txs, []types.TxKey)
+
 	// ReapMaxBytesMaxGas reaps transactions from the mempool up to maxBytes
 	// bytes total with the condition that the total gasWanted must be less than
 	// maxGas and that the total estimated gas used is less than maxGasEstimated.

internal/mempool/types.go

@@ -476,6 +476,10 @@ func (blockExec *BlockExecutor) GetMissingTxs(txKeys []types.TxKey) []types.TxKe
 	return missingTxKeys
 }
 
+func (blockExec *BlockExecutor) SafeGetTxsByKeys(txKeys []types.TxKey) (types.Txs, []types.TxKey) {
+	return blockExec.mempool.SafeGetTxsForKeys(txKeys)
+}
+
 func (blockExec *BlockExecutor) CheckTxFromPeerProposal(ctx context.Context, tx types.Tx) {
 	// Ignore errors from CheckTx because there could be benign errors due to the same tx being
 	// inserted into the mempool from gossiping. Since such simultaneous insertion could result in