From 2c1504f93c6da654e1331ea6a177412a14a50a73 Mon Sep 17 00:00:00 2001
From: Jae Kwon <jae@tendermint.com>
Date: Sun, 6 Mar 2016 15:08:32 -0800
Subject: [PATCH] Reap takes maxTxs

---
 consensus/state.go      | 13 +------------
 mempool/mempool.go      | 15 ++++++++++-----
 mempool/mempool_test.go |  2 +-
 rpc/core/mempool.go     |  2 +-
 4 files changed, 13 insertions(+), 19 deletions(-)

diff --git a/consensus/state.go b/consensus/state.go
index 6b12d05d..51108dab 100644
--- a/consensus/state.go
+++ b/consensus/state.go
@@ -877,19 +877,8 @@ func (cs *ConsensusState) createProposalBlock() (block *types.Block, blockParts
 		return
 	}
 
-	maxBlockSize := config.GetInt("block_size")
-
 	// Mempool validated transactions
-	// if block_size < 0, no txs will be included
-	var txs []types.Tx
-	if maxBlockSize >= 0 {
-		txs = cs.mempool.Reap()
-	}
-
-	// Cap the number of txs in a block
-	if maxBlockSize > 0 && maxBlockSize < len(txs) {
-		txs = txs[:maxBlockSize]
-	}
+	txs := cs.mempool.Reap(config.GetInt("block_size"))
 
 	block = &types.Block{
 		Header: &types.Header{
diff --git a/mempool/mempool.go b/mempool/mempool.go
index 1d0a375d..4385f385 100644
--- a/mempool/mempool.go
+++ b/mempool/mempool.go
@@ -183,7 +183,8 @@ func (mem *Mempool) resCbRecheck(req *tmsp.Request, res *tmsp.Response) {
 }
 
 // Get the valid transactions remaining
-func (mem *Mempool) Reap() []types.Tx {
+// If maxTxs is 0, there is no cap.
+func (mem *Mempool) Reap(maxTxs int) []types.Tx {
 	mem.proxyMtx.Lock()
 	defer mem.proxyMtx.Unlock()
 
@@ -192,13 +193,17 @@ func (mem *Mempool) Reap() []types.Tx {
 		time.Sleep(time.Millisecond * 10)
 	}
 
-	txs := mem.collectTxs()
+	txs := mem.collectTxs(maxTxs)
 	return txs
 }
 
-func (mem *Mempool) collectTxs() []types.Tx {
-	txs := make([]types.Tx, 0, mem.txs.Len())
-	for e := mem.txs.Front(); e != nil; e = e.Next() {
+// maxTxs: 0 means uncapped
+func (mem *Mempool) collectTxs(maxTxs int) []types.Tx {
+	if maxTxs == 0 {
+		maxTxs = mem.txs.Len()
+	}
+	txs := make([]types.Tx, 0, MinInt(mem.txs.Len(), maxTxs))
+	for e := mem.txs.Front(); e != nil && len(txs) < maxTxs; e = e.Next() {
 		memTx := e.Value.(*mempoolTx)
 		txs = append(txs, memTx.tx)
 	}
diff --git a/mempool/mempool_test.go b/mempool/mempool_test.go
index 3fd80bf5..9b57cfdb 100644
--- a/mempool/mempool_test.go
+++ b/mempool/mempool_test.go
@@ -49,7 +49,7 @@ func TestSerialReap(t *testing.T) {
 	}
 
 	reapCheck := func(exp int) {
-		txs := mempool.Reap()
+		txs := mempool.Reap(0)
 		if len(txs) != exp {
 			t.Fatalf("Expected to reap %v txs but got %v", exp, len(txs))
 		}
diff --git a/rpc/core/mempool.go b/rpc/core/mempool.go
index 61bf08a5..e5a34026 100644
--- a/rpc/core/mempool.go
+++ b/rpc/core/mempool.go
@@ -36,6 +36,6 @@ func BroadcastTxSync(tx types.Tx) (*ctypes.ResultBroadcastTx, error) {
 }
 
 func UnconfirmedTxs() (*ctypes.ResultUnconfirmedTxs, error) {
-	txs := mempoolReactor.Mempool.Reap()
+	txs := mempoolReactor.Mempool.Reap(0)
 	return &ctypes.ResultUnconfirmedTxs{len(txs), txs}, nil
 }