chain: share the same bitcoind connection between multiple rescan clients
This commit is contained in:
parent
27e22b1f79
commit
00428d5828
2 changed files with 1001 additions and 592 deletions
File diff suppressed because it is too large
Load diff
362
chain/bitcoind_conn.go
Normal file
362
chain/bitcoind_conn.go
Normal file
|
@ -0,0 +1,362 @@
|
|||
package chain
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"net"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/btcsuite/btcd/chaincfg"
|
||||
"github.com/btcsuite/btcd/chaincfg/chainhash"
|
||||
"github.com/btcsuite/btcd/rpcclient"
|
||||
"github.com/btcsuite/btcd/wire"
|
||||
"github.com/lightninglabs/gozmq"
|
||||
)
|
||||
|
||||
// BitcoindConn represents a persistent client connection to a bitcoind node
|
||||
// that listens for events read from a ZMQ connection.
|
||||
type BitcoindConn struct {
|
||||
started int32 // To be used atomically.
|
||||
stopped int32 // To be used atomically.
|
||||
|
||||
// rescanClientCounter is an atomic counter that assigns a unique ID to
|
||||
// each new bitcoind rescan client using the current bitcoind
|
||||
// connection.
|
||||
rescanClientCounter uint64
|
||||
|
||||
// chainParams identifies the current network the bitcoind node is
|
||||
// running on.
|
||||
chainParams *chaincfg.Params
|
||||
|
||||
// client is the RPC client to the bitcoind node.
|
||||
client *rpcclient.Client
|
||||
|
||||
// zmqBlockHost is the host listening for ZMQ connections that will be
|
||||
// responsible for delivering raw transaction events.
|
||||
zmqBlockHost string
|
||||
|
||||
// zmqTxHost is the host listening for ZMQ connections that will be
|
||||
// responsible for delivering raw transaction events.
|
||||
zmqTxHost string
|
||||
|
||||
// zmqPollInterval is the interval at which we'll attempt to retrieve an
|
||||
// event from the ZMQ connection.
|
||||
zmqPollInterval time.Duration
|
||||
|
||||
// rescanClients is the set of active bitcoind rescan clients to which
|
||||
// ZMQ event notfications will be sent to.
|
||||
rescanClientsMtx sync.Mutex
|
||||
rescanClients map[uint64]*BitcoindClient
|
||||
|
||||
quit chan struct{}
|
||||
wg sync.WaitGroup
|
||||
}
|
||||
|
||||
// NewBitcoindConn creates a client connection to the node described by the host
|
||||
// string. The connection is not established immediately, but must be done using
|
||||
// the Start method. If the remote node does not operate on the same bitcoin
|
||||
// network as described by the passed chain parameters, the connection will be
|
||||
// disconnected.
|
||||
func NewBitcoindConn(chainParams *chaincfg.Params,
|
||||
host, user, pass, zmqBlockHost, zmqTxHost string,
|
||||
zmqPollInterval time.Duration) (*BitcoindConn, error) {
|
||||
|
||||
clientCfg := &rpcclient.ConnConfig{
|
||||
Host: host,
|
||||
User: user,
|
||||
Pass: pass,
|
||||
DisableAutoReconnect: false,
|
||||
DisableConnectOnNew: true,
|
||||
DisableTLS: true,
|
||||
HTTPPostMode: true,
|
||||
}
|
||||
|
||||
client, err := rpcclient.New(clientCfg, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
conn := &BitcoindConn{
|
||||
chainParams: chainParams,
|
||||
client: client,
|
||||
zmqBlockHost: zmqBlockHost,
|
||||
zmqTxHost: zmqTxHost,
|
||||
zmqPollInterval: zmqPollInterval,
|
||||
rescanClients: make(map[uint64]*BitcoindClient),
|
||||
quit: make(chan struct{}),
|
||||
}
|
||||
|
||||
return conn, nil
|
||||
}
|
||||
|
||||
// Start attempts to establish a RPC and ZMQ connection to a bitcoind node. If
|
||||
// successful, a goroutine is spawned to read events from the ZMQ connection.
|
||||
// It's possible for this function to fail due to a limited number of connection
|
||||
// attempts. This is done to prevent waiting forever on the connection to be
|
||||
// established in the case that the node is down.
|
||||
func (c *BitcoindConn) Start() error {
|
||||
if !atomic.CompareAndSwapInt32(&c.started, 0, 1) {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Verify that the node is running on the expected network.
|
||||
net, err := c.getCurrentNet()
|
||||
if err != nil {
|
||||
c.client.Disconnect()
|
||||
return err
|
||||
}
|
||||
if net != c.chainParams.Net {
|
||||
c.client.Disconnect()
|
||||
return fmt.Errorf("expected network %v, got %v",
|
||||
c.chainParams.Net, net)
|
||||
}
|
||||
|
||||
// Establish two different ZMQ connections to bitcoind to retrieve block
|
||||
// and transaction event notifications. We'll use two as a separation of
|
||||
// concern to ensure one type of event isn't dropped from the connection
|
||||
// queue due to another type of event filling it up.
|
||||
zmqBlockConn, err := gozmq.Subscribe(
|
||||
c.zmqBlockHost, []string{"rawblock"}, c.zmqPollInterval,
|
||||
)
|
||||
if err != nil {
|
||||
c.client.Disconnect()
|
||||
return fmt.Errorf("unable to subscribe for zmq block events: "+
|
||||
"%v", err)
|
||||
}
|
||||
|
||||
zmqTxConn, err := gozmq.Subscribe(
|
||||
c.zmqTxHost, []string{"rawtx"}, c.zmqPollInterval,
|
||||
)
|
||||
if err != nil {
|
||||
c.client.Disconnect()
|
||||
return fmt.Errorf("unable to subscribe for zmq tx events: %v",
|
||||
err)
|
||||
}
|
||||
|
||||
c.wg.Add(2)
|
||||
go c.blockEventHandler(zmqBlockConn)
|
||||
go c.txEventHandler(zmqTxConn)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Stop terminates the RPC and ZMQ connection to a bitcoind node and removes any
|
||||
// active rescan clients.
|
||||
func (c *BitcoindConn) Stop() {
|
||||
if !atomic.CompareAndSwapInt32(&c.stopped, 0, 1) {
|
||||
return
|
||||
}
|
||||
|
||||
for _, client := range c.rescanClients {
|
||||
client.Stop()
|
||||
}
|
||||
|
||||
close(c.quit)
|
||||
c.client.Shutdown()
|
||||
|
||||
c.client.WaitForShutdown()
|
||||
c.wg.Wait()
|
||||
}
|
||||
|
||||
// blockEventHandler reads raw blocks events from the ZMQ block socket and
|
||||
// forwards them along to the current rescan clients.
|
||||
//
|
||||
// NOTE: This must be run as a goroutine.
|
||||
func (c *BitcoindConn) blockEventHandler(conn *gozmq.Conn) {
|
||||
defer c.wg.Done()
|
||||
defer conn.Close()
|
||||
|
||||
log.Info("Started listening for bitcoind block notifications via ZMQ ",
|
||||
"on", c.zmqBlockHost)
|
||||
|
||||
for {
|
||||
// Before attempting to read from the ZMQ socket, we'll make
|
||||
// sure to check if we've been requested to shut down.
|
||||
select {
|
||||
case <-c.quit:
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
// Poll an event from the ZMQ socket. It's possible that the
|
||||
// connection to the socket continuously times out, so we'll
|
||||
// prevent logging this error to prevent spamming the logs.
|
||||
msgBytes, err := conn.Receive()
|
||||
if err != nil {
|
||||
err, ok := err.(net.Error)
|
||||
if !ok || !err.Timeout() {
|
||||
log.Error(err)
|
||||
}
|
||||
|
||||
continue
|
||||
}
|
||||
|
||||
// We have an event! We'll now ensure it is a block event,
|
||||
// deserialize it, and report it to the different rescan
|
||||
// clients.
|
||||
eventType := string(msgBytes[0])
|
||||
switch eventType {
|
||||
case "rawblock":
|
||||
block := &wire.MsgBlock{}
|
||||
r := bytes.NewReader(msgBytes[1])
|
||||
if err := block.Deserialize(r); err != nil {
|
||||
log.Errorf("Unable to deserialize block: %v",
|
||||
err)
|
||||
continue
|
||||
}
|
||||
|
||||
c.rescanClientsMtx.Lock()
|
||||
for _, client := range c.rescanClients {
|
||||
select {
|
||||
case client.zmqBlockNtfns <- block:
|
||||
case <-client.quit:
|
||||
case <-c.quit:
|
||||
c.rescanClientsMtx.Unlock()
|
||||
return
|
||||
}
|
||||
}
|
||||
c.rescanClientsMtx.Unlock()
|
||||
default:
|
||||
log.Warnf("Received unexpected event type from "+
|
||||
"rawblock subscription: %v", eventType)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// txEventHandler reads raw blocks events from the ZMQ block socket and forwards
|
||||
// them along to the current rescan clients.
|
||||
//
|
||||
// NOTE: This must be run as a goroutine.
|
||||
func (c *BitcoindConn) txEventHandler(conn *gozmq.Conn) {
|
||||
defer c.wg.Done()
|
||||
defer conn.Close()
|
||||
|
||||
log.Info("Started listening for bitcoind transaction notifications "+
|
||||
"via ZMQ on ", c.zmqTxHost)
|
||||
|
||||
for {
|
||||
// Before attempting to read from the ZMQ socket, we'll make
|
||||
// sure to check if we've been requested to shut down.
|
||||
select {
|
||||
case <-c.quit:
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
// Poll an event from the ZMQ socket. It's possible that the
|
||||
// connection to the socket continuously times out, so we'll
|
||||
// prevent logging this error to prevent spamming the logs.
|
||||
msgBytes, err := conn.Receive()
|
||||
if err != nil {
|
||||
err, ok := err.(net.Error)
|
||||
if !ok || !err.Timeout() {
|
||||
log.Error(err)
|
||||
}
|
||||
|
||||
continue
|
||||
}
|
||||
|
||||
// We have an event! We'll now ensure it is a transaction event,
|
||||
// deserialize it, and report it to the different rescan
|
||||
// clients.
|
||||
eventType := string(msgBytes[0])
|
||||
switch eventType {
|
||||
case "rawtx":
|
||||
tx := &wire.MsgTx{}
|
||||
r := bytes.NewReader(msgBytes[1])
|
||||
if err := tx.Deserialize(r); err != nil {
|
||||
log.Errorf("Unable to deserialize "+
|
||||
"transaction: %v", err)
|
||||
continue
|
||||
}
|
||||
|
||||
c.rescanClientsMtx.Lock()
|
||||
for _, client := range c.rescanClients {
|
||||
select {
|
||||
case client.zmqTxNtfns <- tx:
|
||||
case <-client.quit:
|
||||
case <-c.quit:
|
||||
c.rescanClientsMtx.Unlock()
|
||||
return
|
||||
}
|
||||
}
|
||||
c.rescanClientsMtx.Unlock()
|
||||
default:
|
||||
log.Warnf("Received unexpected event type from rawtx "+
|
||||
"subscription: %v", eventType)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// getCurrentNet returns the network on which the bitcoind node is running.
|
||||
func (c *BitcoindConn) getCurrentNet() (wire.BitcoinNet, error) {
|
||||
hash, err := c.client.GetBlockHash(0)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
switch *hash {
|
||||
case *chaincfg.TestNet3Params.GenesisHash:
|
||||
return chaincfg.TestNet3Params.Net, nil
|
||||
case *chaincfg.RegressionNetParams.GenesisHash:
|
||||
return chaincfg.RegressionNetParams.Net, nil
|
||||
case *chaincfg.MainNetParams.GenesisHash:
|
||||
return chaincfg.MainNetParams.Net, nil
|
||||
default:
|
||||
return 0, fmt.Errorf("unknown network with genesis hash %v", hash)
|
||||
}
|
||||
}
|
||||
|
||||
// NewBitcoindClient returns a bitcoind client using the current bitcoind
|
||||
// connection. This allows us to share the same connection using multiple
|
||||
// clients. The birthday signifies the earliest time for which we should begin
|
||||
// scanning the chain.
|
||||
func (c *BitcoindConn) NewBitcoindClient(birthday time.Time) *BitcoindClient {
|
||||
return &BitcoindClient{
|
||||
quit: make(chan struct{}),
|
||||
|
||||
id: atomic.AddUint64(&c.rescanClientCounter, 1),
|
||||
|
||||
birthday: birthday,
|
||||
chainParams: c.chainParams,
|
||||
chainConn: c,
|
||||
|
||||
rescanUpdate: make(chan interface{}),
|
||||
watchedAddresses: make(map[string]struct{}),
|
||||
watchedOutPoints: make(map[wire.OutPoint]struct{}),
|
||||
watchedTxs: make(map[chainhash.Hash]struct{}),
|
||||
|
||||
notificationQueue: NewConcurrentQueue(20),
|
||||
zmqTxNtfns: make(chan *wire.MsgTx),
|
||||
zmqBlockNtfns: make(chan *wire.MsgBlock),
|
||||
|
||||
mempool: make(map[chainhash.Hash]struct{}),
|
||||
expiredMempool: make(map[int32]map[chainhash.Hash]struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// AddClient adds a client to the set of active rescan clients of the current
|
||||
// chain connection. This allows the connection to include the specified client
|
||||
// in its notification delivery.
|
||||
//
|
||||
// NOTE: This function is safe for concurrent access.
|
||||
func (c *BitcoindConn) AddClient(client *BitcoindClient) {
|
||||
c.rescanClientsMtx.Lock()
|
||||
defer c.rescanClientsMtx.Unlock()
|
||||
|
||||
c.rescanClients[client.id] = client
|
||||
}
|
||||
|
||||
// RemoveClient removes the client with the given ID from the set of active
|
||||
// rescan clients. Once removed, the client will no longer receive block and
|
||||
// transaction notifications from the chain connection.
|
||||
//
|
||||
// NOTE: This function is safe for concurrent access.
|
||||
func (c *BitcoindConn) RemoveClient(id uint64) {
|
||||
c.rescanClientsMtx.Lock()
|
||||
defer c.rescanClientsMtx.Unlock()
|
||||
|
||||
delete(c.rescanClients, id)
|
||||
}
|
Loading…
Reference in a new issue