8 log "github.com/sirupsen/logrus"
10 "github.com/vapor/protocol"
11 "github.com/vapor/protocol/bc"
12 "github.com/vapor/protocol/bc/types"
16 type notificationBlockConnected types.Block
17 type notificationBlockDisconnected types.Block
18 type notificationTxAcceptedByMempool types.Tx
20 // Notification control requests
21 type notificationRegisterClient WSClient
22 type notificationUnregisterClient WSClient
23 type notificationRegisterBlocks WSClient
24 type notificationUnregisterBlocks WSClient
25 type notificationRegisterNewMempoolTxs WSClient
26 type notificationUnregisterNewMempoolTxs WSClient
28 // NotificationType represents the type of a notification message.
29 type NotificationType int
31 // Constants for the type of a notification message.
33 // NTBlockConnected indicates the associated block was connected to the main chain.
34 NTRawBlockConnected NotificationType = iota
35 // NTBlockDisconnected indicates the associated block was disconnected from the main chain.
36 NTRawBlockDisconnected
41 // notificationTypeStrings is a map of notification types back to their constant
42 // names for pretty printing.
43 var notificationTypeStrings = map[NotificationType]string{
44 NTRawBlockConnected: "raw_blocks_connected",
45 NTRawBlockDisconnected: "raw_blocks_disconnected",
46 NTNewTransaction: "new_transaction",
47 NTRequestStatus: "request_status",
50 // String returns the NotificationType in human-readable form.
51 func (n NotificationType) String() string {
52 if s, ok := notificationTypeStrings[n]; ok {
55 return fmt.Sprintf("Unknown Notification Type (%d)", int(n))
58 type statusInfo struct {
63 // WSNotificationManager is a connection and notification manager used for
64 // websockets. It allows websocket clients to register for notifications they
65 // are interested in. When an event happens elsewhere in the code such as
66 // transactions being added to the memory pool or block connects/disconnects,
67 // the notification manager is provided with the relevant details needed to
68 // figure out which websocket clients need to be notified based on what they
69 // have registered for and notifies them accordingly. It is also used to keep
70 // track of all connected websocket clients.
71 type WSNotificationManager struct {
72 // queueNotification queues a notification for handling.
73 queueNotification chan interface{}
75 // notificationMsgs feeds notificationHandler with notifications
76 // and client (un)registeration requests from a queue as well as
77 // registeration and unregisteration requests from clients.
78 notificationMsgs chan interface{}
80 // Access channel for current number of connected clients.
87 maxNumConcurrentReqs int
92 // NewWsNotificationManager returns a new notification manager ready for use. See WSNotificationManager for more details.
93 func NewWsNotificationManager(maxNumWebsockets int, maxNumConcurrentReqs int, chain *protocol.Chain) *WSNotificationManager {
96 header := chain.BestBlockHeader()
97 status.BestHeight = header.Height
98 status.BestHash = header.Hash()
100 return &WSNotificationManager{
101 queueNotification: make(chan interface{}),
102 notificationMsgs: make(chan interface{}),
103 numClients: make(chan int),
104 quit: make(chan struct{}),
105 MaxNumWebsockets: maxNumWebsockets,
106 maxNumConcurrentReqs: maxNumConcurrentReqs,
112 // queueHandler manages a queue of empty interfaces, reading from in and
113 // sending the oldest unsent to out. This handler stops when either of the
114 // in or quit channels are closed, and closes out before returning, without
115 // waiting to send any variables still remaining in the queue.
116 func queueHandler(in <-chan interface{}, out chan<- interface{}, quit <-chan struct{}) {
120 dequeue chan<- interface{}
130 // Sender closed input channel.
134 // Either send to out immediately if skipQueue is
135 // non-nil (queue is empty) and reader is ready,
136 // or append to the queue and send later.
147 case dequeue <- next:
149 q[len(q)-1] = nil // avoid leak
165 func (m *WSNotificationManager) sendNotification(typ NotificationType, data interface{}) {
167 case NTRawBlockConnected:
168 block, ok := data.(*types.Block)
170 log.WithField("module", logModule).Error("Chain connected notification is not a block")
174 // Notify registered websocket clients of incoming block.
175 m.NotifyBlockConnected(block)
177 case NTRawBlockDisconnected:
178 block, ok := data.(*types.Block)
180 log.WithField("module", logModule).Error("Chain disconnected notification is not a block")
184 // Notify registered websocket clients.
185 m.NotifyBlockDisconnected(block)
189 // queueHandler maintains a queue of notifications and notification handler
191 func (m *WSNotificationManager) queueHandler() {
192 queueHandler(m.queueNotification, m.notificationMsgs, m.quit)
196 // NotifyBlockConnected passes a block newly-connected to the best chain
197 // to the notification manager for block and transaction notification
199 func (m *WSNotificationManager) NotifyBlockConnected(block *types.Block) {
201 case m.queueNotification <- (*notificationBlockConnected)(block):
206 // NotifyBlockDisconnected passes a block disconnected from the best chain
207 // to the notification manager for block notification processing.
208 func (m *WSNotificationManager) NotifyBlockDisconnected(block *types.Block) {
210 case m.queueNotification <- (*notificationBlockDisconnected)(block):
215 // NotifyMempoolTx passes a transaction accepted by mempool to the
216 // notification manager for transaction notification processing. If
217 // isNew is true, the tx is is a new transaction, rather than one
218 // added to the mempool during a reorg.
219 func (m *WSNotificationManager) NotifyMempoolTx(tx *types.Tx) {
221 case m.queueNotification <- (*notificationTxAcceptedByMempool)(tx):
226 // notificationHandler reads notifications and control messages from the queue handler and processes one at a time.
227 func (m *WSNotificationManager) notificationHandler() {
228 // clients is a map of all currently connected websocket clients.
229 clients := make(map[chan struct{}]*WSClient)
230 blockNotifications := make(map[chan struct{}]*WSClient)
231 txNotifications := make(map[chan struct{}]*WSClient)
236 case n, ok := <-m.notificationMsgs:
240 switch n := n.(type) {
241 case *notificationBlockConnected:
242 block := (*types.Block)(n)
243 if len(blockNotifications) != 0 {
244 m.notifyBlockConnected(blockNotifications, block)
247 case *notificationBlockDisconnected:
248 block := (*types.Block)(n)
249 if len(blockNotifications) != 0 {
250 m.notifyBlockDisconnected(blockNotifications, block)
253 case *notificationTxAcceptedByMempool:
255 if len(txNotifications) != 0 {
256 m.notifyForNewTx(txNotifications, tx)
259 case *notificationRegisterBlocks:
260 wsc := (*WSClient)(n)
261 blockNotifications[wsc.quit] = wsc
263 case *notificationUnregisterBlocks:
264 wsc := (*WSClient)(n)
265 delete(blockNotifications, wsc.quit)
267 case *notificationRegisterNewMempoolTxs:
268 wsc := (*WSClient)(n)
269 txNotifications[wsc.quit] = wsc
271 case *notificationUnregisterNewMempoolTxs:
272 wsc := (*WSClient)(n)
273 delete(txNotifications, wsc.quit)
275 case *notificationRegisterClient:
276 wsc := (*WSClient)(n)
277 clients[wsc.quit] = wsc
279 case *notificationUnregisterClient:
280 wsc := (*WSClient)(n)
281 delete(blockNotifications, wsc.quit)
282 delete(txNotifications, wsc.quit)
283 delete(clients, wsc.quit)
286 log.Warnf("Unhandled notification type")
289 case m.numClients <- len(clients):
296 for _, c := range clients {
302 // NumClients returns the number of clients actively being served.
303 func (m *WSNotificationManager) NumClients() (n int) {
305 case n = <-m.numClients:
311 // IsMaxConnect returns whether the maximum connection is exceeded
312 func (m *WSNotificationManager) IsMaxConnect() bool {
313 return m.NumClients() >= m.MaxNumWebsockets
316 // RegisterBlockUpdates requests block update notifications to the passed websocket client.
317 func (m *WSNotificationManager) RegisterBlockUpdates(wsc *WSClient) {
318 m.queueNotification <- (*notificationRegisterBlocks)(wsc)
321 // UnregisterBlockUpdates removes block update notifications for the passed websocket client.
322 func (m *WSNotificationManager) UnregisterBlockUpdates(wsc *WSClient) {
323 m.queueNotification <- (*notificationUnregisterBlocks)(wsc)
326 // notifyBlockConnected notifies websocket clients that have registered for block updates when a block is connected to the main chain.
327 func (*WSNotificationManager) notifyBlockConnected(clients map[chan struct{}]*WSClient, block *types.Block) {
328 resp := NewWSResponse(NTRawBlockConnected.String(), block, nil)
329 marshalledJSON, err := json.Marshal(resp)
331 log.WithFields(log.Fields{"module": logModule, "error": err}).Error("Failed to marshal block connected notification")
335 for _, wsc := range clients {
336 wsc.QueueNotification(marshalledJSON)
340 // notifyBlockDisconnected notifies websocket clients that have registered for block updates
341 // when a block is disconnected from the main chain (due to a reorganize).
342 func (*WSNotificationManager) notifyBlockDisconnected(clients map[chan struct{}]*WSClient, block *types.Block) {
343 resp := NewWSResponse(NTRawBlockDisconnected.String(), block, nil)
344 marshalledJSON, err := json.Marshal(resp)
346 log.WithField("error", err).Error("Failed to marshal block Disconnected notification")
350 for _, wsc := range clients {
351 wsc.QueueNotification(marshalledJSON)
355 // RegisterNewMempoolTxsUpdates requests notifications to the passed websocket
356 // client when new transactions are added to the memory pool.
357 func (m *WSNotificationManager) RegisterNewMempoolTxsUpdates(wsc *WSClient) {
358 m.queueNotification <- (*notificationRegisterNewMempoolTxs)(wsc)
361 // UnregisterNewMempoolTxsUpdates removes notifications to the passed websocket
362 // client when new transaction are added to the memory pool.
363 func (m *WSNotificationManager) UnregisterNewMempoolTxsUpdates(wsc *WSClient) {
364 m.queueNotification <- (*notificationUnregisterNewMempoolTxs)(wsc)
367 // notifyForNewTx notifies websocket clients that have registered for updates
368 // when a new transaction is added to the memory pool.
369 func (m *WSNotificationManager) notifyForNewTx(clients map[chan struct{}]*WSClient, tx *types.Tx) {
370 resp := NewWSResponse(NTNewTransaction.String(), tx, nil)
371 marshalledJSON, err := json.Marshal(resp)
373 log.WithFields(log.Fields{"module": logModule, "error": err}).Error("Failed to marshal tx notification")
377 for _, wsc := range clients {
378 wsc.QueueNotification(marshalledJSON)
382 // AddClient adds the passed websocket client to the notification manager.
383 func (m *WSNotificationManager) AddClient(wsc *WSClient) {
384 m.queueNotification <- (*notificationRegisterClient)(wsc)
387 // RemoveClient removes the passed websocket client and all notifications registered for it.
388 func (m *WSNotificationManager) RemoveClient(wsc *WSClient) {
390 case m.queueNotification <- (*notificationUnregisterClient)(wsc):
395 func (m *WSNotificationManager) blockNotify() {
404 for !m.chain.InMainChain(m.status.BestHash) {
405 block, err := m.chain.GetBlockByHash(&m.status.BestHash)
407 log.WithFields(log.Fields{"module": logModule, "err": err}).Error("blockNotify GetBlockByHash")
410 m.updateStatus(block)
411 m.sendNotification(NTRawBlockDisconnected, block)
414 block, _ := m.chain.GetBlockByHeight(m.status.BestHeight + 1)
420 if m.status.BestHash != block.PreviousBlockHash {
421 log.WithFields(log.Fields{"module": logModule, "blockHeight": block.Height, "previousBlockHash": m.status.BestHash, "rcvBlockPrevHash": block.PreviousBlockHash}).Warning("The previousBlockHash of the received block is not the same as the hash of the previous block")
425 m.updateStatus(block)
426 m.sendNotification(NTRawBlockConnected, block)
431 func (m *WSNotificationManager) blockWaiter() {
433 case <-m.chain.BlockWaiter(m.status.BestHeight + 1):
438 func (m *WSNotificationManager) updateStatus(block *types.Block) {
439 m.status.BestHeight = block.Height
440 m.status.BestHash = block.Hash()
443 // Start starts the goroutines required for the manager to queue and process websocket client notifications.
444 func (m *WSNotificationManager) Start() {
448 go m.notificationHandler()
451 // WaitForShutdown blocks until all notification manager goroutines have finished.
452 func (m *WSNotificationManager) WaitForShutdown() {
456 // Shutdown shuts down the manager, stopping the notification queue and notification handler goroutines.
457 func (m *WSNotificationManager) Shutdown() {