github.com/cryptohub-digital/blockbook@v0.3.5-0.20240403155730-99ab40b9104c/bchain/mq.go (about)

     1  package bchain
     2  
     3  import (
     4  	"context"
     5  	"encoding/binary"
     6  	"time"
     7  
     8  	"github.com/golang/glog"
     9  	zmq "github.com/pebbe/zmq4"
    10  )
    11  
    12  // MQ is message queue listener handle
    13  type MQ struct {
    14  	context   *zmq.Context
    15  	socket    *zmq.Socket
    16  	isRunning bool
    17  	finished  chan error
    18  	binding   string
    19  }
    20  
    21  // NotificationType is type of notification
    22  type NotificationType int
    23  
    24  const (
    25  	// NotificationUnknown is unknown
    26  	NotificationUnknown NotificationType = iota
    27  	// NotificationNewBlock message is sent when there is a new block to be imported
    28  	NotificationNewBlock NotificationType = iota
    29  	// NotificationNewTx message is sent when there is a new mempool transaction
    30  	NotificationNewTx NotificationType = iota
    31  )
    32  
    33  // NewMQ creates new Bitcoind ZeroMQ listener
    34  // callback function receives messages
    35  func NewMQ(binding string, callback func(NotificationType)) (*MQ, error) {
    36  	context, err := zmq.NewContext()
    37  	if err != nil {
    38  		return nil, err
    39  	}
    40  	socket, err := context.NewSocket(zmq.SUB)
    41  	if err != nil {
    42  		return nil, err
    43  	}
    44  	err = socket.SetSubscribe("hashblock")
    45  	if err != nil {
    46  		return nil, err
    47  	}
    48  	err = socket.SetSubscribe("hashtx")
    49  	if err != nil {
    50  		return nil, err
    51  	}
    52  	// for now do not use raw subscriptions - we would have to handle skipped/lost notifications from zeromq
    53  	// on each notification we do sync or syncmempool respectively
    54  	// socket.SetSubscribe("rawblock")
    55  	// socket.SetSubscribe("rawtx")
    56  	err = socket.Connect(binding)
    57  	if err != nil {
    58  		return nil, err
    59  	}
    60  	glog.Info("MQ listening to ", binding)
    61  	mq := &MQ{context, socket, true, make(chan error), binding}
    62  	go mq.run(callback)
    63  	return mq, nil
    64  }
    65  
    66  func (mq *MQ) run(callback func(NotificationType)) {
    67  	defer func() {
    68  		if r := recover(); r != nil {
    69  			glog.Error("MQ loop recovered from ", r)
    70  		}
    71  		mq.isRunning = false
    72  		glog.Info("MQ loop terminated")
    73  		mq.finished <- nil
    74  	}()
    75  	mq.isRunning = true
    76  	repeatedError := false
    77  	for {
    78  		msg, err := mq.socket.RecvMessageBytes(0)
    79  		if err != nil {
    80  			if zmq.AsErrno(err) == zmq.Errno(zmq.ETERM) || err.Error() == "Socket is closed" {
    81  				break
    82  			}
    83  			// suppress logging of error for the first time
    84  			// programs built with Go 1.14 will receive more signals
    85  			// the error should be resolved by retrying the call
    86  			// see https://golang.org/doc/go1.14#runtime
    87  			if repeatedError {
    88  				glog.Error("MQ RecvMessageBytes error ", err, ", ", zmq.AsErrno(err))
    89  			}
    90  			repeatedError = true
    91  			time.Sleep(100 * time.Millisecond)
    92  		} else {
    93  			repeatedError = false
    94  		}
    95  		if len(msg) >= 3 {
    96  			var nt NotificationType
    97  			switch string(msg[0]) {
    98  			case "hashblock":
    99  				nt = NotificationNewBlock
   100  			case "hashtx":
   101  				nt = NotificationNewTx
   102  			default:
   103  				nt = NotificationUnknown
   104  				glog.Infof("MQ: NotificationUnknown %v", string(msg[0]))
   105  			}
   106  			if glog.V(2) {
   107  				sequence := uint32(0)
   108  				if len(msg[len(msg)-1]) == 4 {
   109  					sequence = binary.LittleEndian.Uint32(msg[len(msg)-1])
   110  				}
   111  				glog.Infof("MQ: %v %s-%d", nt, string(msg[0]), sequence)
   112  			}
   113  			callback(nt)
   114  		}
   115  	}
   116  }
   117  
   118  // Shutdown stops listening to the ZeroMQ and closes the connection
   119  func (mq *MQ) Shutdown(ctx context.Context) error {
   120  	glog.Info("MQ server shutdown")
   121  	if mq.isRunning {
   122  		go func() {
   123  			// if errors in the closing sequence, let it close ungracefully
   124  			if err := mq.socket.SetUnsubscribe("hashtx"); err != nil {
   125  				mq.finished <- err
   126  				return
   127  			}
   128  			if err := mq.socket.SetUnsubscribe("hashblock"); err != nil {
   129  				mq.finished <- err
   130  				return
   131  			}
   132  			if err := mq.socket.Unbind(mq.binding); err != nil {
   133  				mq.finished <- err
   134  				return
   135  			}
   136  			if err := mq.socket.Close(); err != nil {
   137  				mq.finished <- err
   138  				return
   139  			}
   140  			if err := mq.context.Term(); err != nil {
   141  				mq.finished <- err
   142  				return
   143  			}
   144  		}()
   145  		var err error
   146  		select {
   147  		case <-ctx.Done():
   148  			err = ctx.Err()
   149  		case err = <-mq.finished:
   150  		}
   151  		if err != nil {
   152  			return err
   153  		}
   154  	}
   155  	glog.Info("MQ server shutdown finished")
   156  	return nil
   157  }