From abb1b8b388535c5e55acfa5452d61b41d149f771 Mon Sep 17 00:00:00 2001
From: Jonathan Moody <103143855+moodyjon@users.noreply.github.com>
Date: Thu, 14 Jul 2022 17:58:24 -0400
Subject: [PATCH] [rpc mempool] Add support for unbroadcastcount to RPC
 getmempoolinfo.

---
 btcjson/chainsvrresults.go |  1 +
 mempool/mempool.go         | 18 ++++++++++++++++++
 rpcserver.go               |  1 +
 rpcserverhelp.go           |  1 +
 server.go                  |  5 +++++
 5 files changed, 26 insertions(+)

diff --git a/btcjson/chainsvrresults.go b/btcjson/chainsvrresults.go
index 261cdcf4..ee588fbf 100644
--- a/btcjson/chainsvrresults.go
+++ b/btcjson/chainsvrresults.go
@@ -343,6 +343,7 @@ type GetMempoolInfoResult struct {
 	TotalFee         float64 `json:"total_fee"`        // Total fees for the mempool in BTC, ignoring modified fees through prioritizetransaction
 	MemPoolMinFee    float64 `json:"mempoolminfee"`    // Minimum fee rate in BTC/kvB for tx to be accepted. Is the maximum of minrelaytxfee and minimum mempool fee
 	MinRelayTxFee    float64 `json:"minrelaytxfee"`    // Current minimum relay fee for transactions
+	UnbroadcastCount int64   `json:"unbroadcastcount"` // Current number of transactions that haven't passed initial broadcast yet
 }
 
 // NetworksResult models the networks data from the getnetworkinfo command.
diff --git a/mempool/mempool.go b/mempool/mempool.go
index 9f2e4f28..688e532d 100644
--- a/mempool/mempool.go
+++ b/mempool/mempool.go
@@ -235,6 +235,9 @@ type TxPool struct {
 
 	// stats are aggregated over pool, orphans, etc.
 	stats aggregateInfo
+
+	// unbroadcast is a set of transactions yet to be broadcast.
+	unbroadcast map[chainhash.Hash]bool
 }
 
 // Ensure the TxPool type implements the mining.TxSource interface.
@@ -1555,10 +1558,23 @@ func (mp *TxPool) MiningDescs() []*mining.TxDesc {
 	return descs
 }
 
+func (mp *TxPool) AddUnbroadcastTx(hash *chainhash.Hash) {
+	mp.mtx.Lock()
+	mp.unbroadcast[*hash] = true
+	mp.mtx.Unlock()
+}
+
+func (mp *TxPool) RemoveUnbroadcastTx(hash *chainhash.Hash) {
+	mp.mtx.Lock()
+	delete(mp.unbroadcast, *hash)
+	mp.mtx.Unlock()
+}
+
 func (mp *TxPool) MempoolInfo() *btcjson.GetMempoolInfoResult {
 	mp.mtx.RLock()
 	policy := mp.cfg.Policy
 	stats := mp.stats
+	unbroadcastCount := int64(len(mp.unbroadcast))
 	mp.mtx.RUnlock()
 
 	ret := &btcjson.GetMempoolInfoResult{
@@ -1568,6 +1584,7 @@ func (mp *TxPool) MempoolInfo() *btcjson.GetMempoolInfoResult {
 		TotalFee:         btcutil.Amount(stats.totalFee).ToBTC(),
 		MemPoolMinFee:    btcutil.Amount(calcMinRequiredTxRelayFee(1000, policy.MinRelayTxFee)).ToBTC(),
 		MinRelayTxFee:    policy.MinRelayTxFee.ToBTC(),
+		UnbroadcastCount: unbroadcastCount,
 	}
 
 	return ret
@@ -1646,5 +1663,6 @@ func New(cfg *Config) *TxPool {
 		orphansByPrev:  make(map[wire.OutPoint]map[chainhash.Hash]*btcutil.Tx),
 		nextExpireScan: time.Now().Add(orphanExpireScanInterval),
 		outpoints:      make(map[wire.OutPoint]*btcutil.Tx),
+		unbroadcast:    make(map[chainhash.Hash]bool),
 	}
 }
diff --git a/rpcserver.go b/rpcserver.go
index e2f21981..1fc32a06 100644
--- a/rpcserver.go
+++ b/rpcserver.go
@@ -3746,6 +3746,7 @@ func handleSendRawTransaction(s *rpcServer, cmd interface{}, closeChan <-chan st
 	// Keep track of all the sendrawtransaction request txns so that they
 	// can be rebroadcast if they don't make their way into a block.
 	txD := acceptedTxs[0]
+	s.cfg.TxMemPool.AddUnbroadcastTx(txD.Tx.Hash())
 	iv := wire.NewInvVect(wire.InvTypeTx, txD.Tx.Hash())
 	s.cfg.ConnMgr.AddRebroadcastInventory(iv, txD)
 
diff --git a/rpcserverhelp.go b/rpcserverhelp.go
index 8c339855..3a6eb87b 100644
--- a/rpcserverhelp.go
+++ b/rpcserverhelp.go
@@ -459,6 +459,7 @@ var helpDescsEnUS = map[string]string{
 	"getmempoolinforesult-total_fee":        "Total fees for the mempool in LBC, ignoring modified fees through prioritizetransaction",
 	"getmempoolinforesult-mempoolminfee":    "Minimum fee rate in LBC/kvB for tx to be accepted. Is the maximum of minrelaytxfee and minimum mempool fee",
 	"getmempoolinforesult-minrelaytxfee":    "Current minimum relay fee for transactions",
+	"getmempoolinforesult-unbroadcastcount": "Current number of transactions that haven't passed initial broadcast yet",
 
 	// GetMiningInfoResult help.
 	"getmininginforesult-blocks":             "Height of the latest best block",
diff --git a/server.go b/server.go
index 82da9ee1..db13ed64 100644
--- a/server.go
+++ b/server.go
@@ -699,6 +699,11 @@ func (sp *serverPeer) OnGetData(_ *peer.Peer, msg *wire.MsgGetData) {
 			if i == len(msg.InvList)-1 && c != nil {
 				<-c
 			}
+		} else if iv.Type == wire.InvTypeWitnessTx || iv.Type == wire.InvTypeTx {
+			// We interpret fulfilling a GETDATA for a transaction as a
+			// successful initial broadcast and remove it from our
+			// unbroadcast set.
+			sp.server.txMemPool.RemoveUnbroadcastTx(&iv.Hash)
 		}
 		numAdded++
 		waitChan = c