github.com/aychain/blockbook@v0.1.1-0.20181121092459-6d1fc7e07c5b/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  	for {
    77  		msg, err := mq.socket.RecvMessageBytes(0)
    78  		if err != nil {
    79  			if zmq.AsErrno(err) == zmq.Errno(zmq.ETERM) || err.Error() == "Socket is closed" {
    80  				break
    81  			}
    82  			glog.Error("MQ RecvMessageBytes error ", err, ", ", zmq.AsErrno(err))
    83  			time.Sleep(100 * time.Millisecond)
    84  		}
    85  		if msg != nil && len(msg) >= 3 {
    86  			var nt NotificationType
    87  			switch string(msg[0]) {
    88  			case "hashblock":
    89  				nt = NotificationNewBlock
    90  				break
    91  			case "hashtx":
    92  				nt = NotificationNewTx
    93  				break
    94  			default:
    95  				nt = NotificationUnknown
    96  				glog.Infof("MQ: NotificationUnknown %v", string(msg[0]))
    97  			}
    98  			if glog.V(2) {
    99  				sequence := uint32(0)
   100  				if len(msg[len(msg)-1]) == 4 {
   101  					sequence = binary.LittleEndian.Uint32(msg[len(msg)-1])
   102  				}
   103  				glog.Infof("MQ: %v %s-%d", nt, string(msg[0]), sequence)
   104  			}
   105  			callback(nt)
   106  		}
   107  	}
   108  }
   109  
   110  // Shutdown stops listening to the ZeroMQ and closes the connection
   111  func (mq *MQ) Shutdown(ctx context.Context) error {
   112  	glog.Info("MQ server shutdown")
   113  	if mq.isRunning {
   114  		go func() {
   115  			// if errors in the closing sequence, let it close ungracefully
   116  			if err := mq.socket.SetUnsubscribe("hashtx"); err != nil {
   117  				mq.finished <- err
   118  				return
   119  			}
   120  			if err := mq.socket.SetUnsubscribe("hashblock"); err != nil {
   121  				mq.finished <- err
   122  				return
   123  			}
   124  			if err := mq.socket.Unbind(mq.binding); err != nil {
   125  				mq.finished <- err
   126  				return
   127  			}
   128  			if err := mq.socket.Close(); err != nil {
   129  				mq.finished <- err
   130  				return
   131  			}
   132  			if err := mq.context.Term(); err != nil {
   133  				mq.finished <- err
   134  				return
   135  			}
   136  		}()
   137  		var err error
   138  		select {
   139  		case <-ctx.Done():
   140  			err = ctx.Err()
   141  		case err = <-mq.finished:
   142  		}
   143  		if err != nil {
   144  			return err
   145  		}
   146  	}
   147  	glog.Info("MQ server shutdown finished")
   148  	return nil
   149  }