github.com/nspcc-dev/neo-go@v0.105.2-0.20240517133400-6be757af3eba/pkg/core/mempool/subscriptions.go (about)

     1  package mempool
     2  
     3  import "github.com/nspcc-dev/neo-go/pkg/core/mempoolevent"
     4  
     5  // RunSubscriptions runs subscriptions goroutine if mempool subscriptions are enabled.
     6  // You should manually free the resources by calling StopSubscriptions on mempool shutdown.
     7  func (mp *Pool) RunSubscriptions() {
     8  	if !mp.subscriptionsEnabled {
     9  		panic("subscriptions are disabled")
    10  	}
    11  	if !mp.subscriptionsOn.Load() {
    12  		mp.subscriptionsOn.Store(true)
    13  		go mp.notificationDispatcher()
    14  	}
    15  }
    16  
    17  // StopSubscriptions stops mempool events loop.
    18  func (mp *Pool) StopSubscriptions() {
    19  	if !mp.subscriptionsEnabled {
    20  		panic("subscriptions are disabled")
    21  	}
    22  	if mp.subscriptionsOn.Load() {
    23  		mp.subscriptionsOn.Store(false)
    24  		close(mp.stopCh)
    25  	}
    26  }
    27  
    28  // SubscribeForTransactions adds the given channel to the new mempool event broadcasting, so when
    29  // there is a new transactions added to the mempool or an existing transaction removed from
    30  // the mempool, you'll receive it via this channel. Make sure you're not changing the received
    31  // mempool events, as it may affect the functionality of other subscribers.
    32  func (mp *Pool) SubscribeForTransactions(ch chan<- mempoolevent.Event) {
    33  	if mp.subscriptionsOn.Load() {
    34  		mp.subCh <- ch
    35  	}
    36  }
    37  
    38  // UnsubscribeFromTransactions unsubscribes the given channel from new mempool notifications,
    39  // you can close it afterwards. Passing non-subscribed channel is a no-op.
    40  func (mp *Pool) UnsubscribeFromTransactions(ch chan<- mempoolevent.Event) {
    41  	if mp.subscriptionsOn.Load() {
    42  		mp.unsubCh <- ch
    43  	}
    44  }
    45  
    46  // notificationDispatcher manages subscription to events and broadcasts new events.
    47  func (mp *Pool) notificationDispatcher() {
    48  	var (
    49  		// These are just sets of subscribers, though modelled as maps
    50  		// for ease of management (not a lot of subscriptions is really
    51  		// expected, but maps are convenient for adding/deleting elements).
    52  		txFeed = make(map[chan<- mempoolevent.Event]bool)
    53  	)
    54  	for {
    55  		select {
    56  		case <-mp.stopCh:
    57  			return
    58  		case sub := <-mp.subCh:
    59  			txFeed[sub] = true
    60  		case unsub := <-mp.unsubCh:
    61  			delete(txFeed, unsub)
    62  		case event := <-mp.events:
    63  			for ch := range txFeed {
    64  				ch <- event
    65  			}
    66  		}
    67  	}
    68  }