lightwalletd/frontend/mempool.go

package frontend

import (
	"encoding/hex"
	"encoding/json"
	"sync"
	"time"

	"github.com/btcsuite/btcd/rpcclient"
	"github.com/sirupsen/logrus"

	"git.hush.is/hush/lightwalletd/common"
	"git.hush.is/hush/lightwalletd/walletrpc"
)

// mempoolMonitor polls the node's mempool ONCE (regardless of how many wallets are streaming) and
// fans each new transaction out to every subscribed GetMempoolStream handler. This replaces
// per-connection polling: N connected wallets no longer each hit getrawmempool/getrawtransaction.
//
// Lifecycle matches the semantics the client's monitor loop expects: while a block is current the
// monitor emits each mempool tx to subscribers exactly once; when a new block is mined it resets
// and closes all subscriber channels, so each handler returns (closing its stream) and the client
// reconnects after re-syncing the block. The poller starts lazily on the first subscription and,
// when no clients are connected, idles without touching the node.
type mempoolMonitor struct {
	mu      sync.Mutex
	subs    map[int]chan *walletrpc.RawTransaction
	nextID  int
	seen    map[string]bool             // txids already emitted for the current block
	ordered []*walletrpc.RawTransaction // emitted txs in arrival order, replayed to late subscribers
}

var (
	sharedMempoolMonitor *mempoolMonitor
	mempoolMonitorOnce   sync.Once
)

// getMempoolMonitor returns the process-wide mempool monitor, starting its poller on first use.
// The client/cache/log are bound once (they are process singletons on the SqlStreamer).
func getMempoolMonitor(client *rpcclient.Client, cache *common.BlockCache, log *logrus.Entry) *mempoolMonitor {
	mempoolMonitorOnce.Do(func() {
		sharedMempoolMonitor = &mempoolMonitor{
			subs: make(map[int]chan *walletrpc.RawTransaction),
			seen: make(map[string]bool),
		}
		go sharedMempoolMonitor.run(client, cache, log)
	})
	return sharedMempoolMonitor
}

// subscribe registers a stream. It returns the subscriber id, a channel of subsequent mempool txs,
// and a snapshot of txs already emitted this block (to be sent first). Registration and snapshot
// are taken atomically, so every tx reaches a subscriber exactly once (snapshot xor channel).
func (m *mempoolMonitor) subscribe() (int, chan *walletrpc.RawTransaction, []*walletrpc.RawTransaction) {
	m.mu.Lock()
	defer m.mu.Unlock()
	id := m.nextID
	m.nextID++
	ch := make(chan *walletrpc.RawTransaction, 256)
	m.subs[id] = ch
	snapshot := make([]*walletrpc.RawTransaction, len(m.ordered))
	copy(snapshot, m.ordered)
	return id, ch, snapshot
}

// unsubscribe removes a stream. Safe to call after a new-block reset already dropped it (the map
// lookup guards against a double close).
func (m *mempoolMonitor) unsubscribe(id int) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if ch, ok := m.subs[id]; ok {
		delete(m.subs, id)
		close(ch)
	}
}

// reset is called when a new block is mined: forget this block's txs and close every subscriber
// stream so clients re-sync the block and reconnect.
func (m *mempoolMonitor) reset() {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.seen = make(map[string]bool)
	m.ordered = nil
	for id, ch := range m.subs {
		delete(m.subs, id)
		close(ch)
	}
}

func (m *mempoolMonitor) run(client *rpcclient.Client, cache *common.BlockCache, log *logrus.Entry) {
	lastHeight := cache.GetLatestBlock()
	for {
		// Track block height even while idle, so a new subscriber isn't immediately closed by a
		// height that advanced during the idle period.
		if h := cache.GetLatestBlock(); h != lastHeight {
			lastHeight = h
			m.reset()
		}

		m.mu.Lock()
		idle := len(m.subs) == 0
		m.mu.Unlock()
		if idle {
			// No one is listening; don't touch the node until a client subscribes.
			time.Sleep(2 * time.Second)
			continue
		}

		height := uint64(lastHeight + 1)

		// List the current mempool txids: getrawmempool(false) -> ["txid", ...].
		mpParams := []json.RawMessage{json.RawMessage("false")}
		result, rpcErr := client.RawRequest("getrawmempool", mpParams)
		if rpcErr != nil {
			log.Warningf("mempool monitor: getrawmempool failed: %s", rpcErr.Error())
			time.Sleep(2 * time.Second)
			continue
		}
		var txids []string
		if err := json.Unmarshal(result, &txids); err != nil {
			log.Warningf("mempool monitor: cannot parse getrawmempool: %s", err.Error())
			time.Sleep(2 * time.Second)
			continue
		}

		for _, txid := range txids {
			m.mu.Lock()
			already := m.seen[txid]
			m.mu.Unlock()
			if already {
				continue
			}

			// Fetch the full serialized tx: getrawtransaction("txid") -> hex. Full bytes are
			// required so shielded memos survive (a CompactTx's 52-byte prefix would not carry a
			// memo). The mempool txid is already in RPC display byte order, so (unlike
			// GetTransaction) no reversal is needed.
			txParams := []json.RawMessage{json.RawMessage("\"" + txid + "\"")}
			txResult, txErr := client.RawRequest("getrawtransaction", txParams)
			if txErr != nil {
				// The tx may have been mined or evicted between listing and fetch; skip it.
				continue
			}
			var txhex string
			if err := json.Unmarshal(txResult, &txhex); err != nil {
				continue
			}
			txBytes, err := hex.DecodeString(txhex)
			if err != nil {
				continue
			}
			rtx := &walletrpc.RawTransaction{Data: txBytes, Height: height}

			// Record + fan out under the lock. The send is non-blocking (default case), so a slow
			// subscriber never stalls the poller or the other subscribers.
			m.mu.Lock()
			if !m.seen[txid] {
				m.seen[txid] = true
				m.ordered = append(m.ordered, rtx)
				for _, ch := range m.subs {
					select {
					case ch <- rtx:
					default:
						// Slow subscriber; drop. Block sync will deliver the confirmed tx later.
					}
				}
			}
			m.mu.Unlock()
		}

		time.Sleep(2 * time.Second)
	}
}