github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/network/p2p/subscription/subscriptionManager.go (about)

     1  package subscription
     2  
     3  import (
     4  	"fmt"
     5  	"sync"
     6  
     7  	"github.com/onflow/flow-go/network"
     8  	"github.com/onflow/flow-go/network/channels"
     9  )
    10  
    11  // ChannelSubscriptionManager manages subscriptions of engines running on the node to channels.
    12  // Each channel should be taken by at most a single engine.
    13  type ChannelSubscriptionManager struct {
    14  	mu              sync.RWMutex
    15  	engines         map[channels.Channel]network.MessageProcessor
    16  	networkUnderlay network.Underlay // the Underlay interface of the network layer
    17  }
    18  
    19  // NewChannelSubscriptionManager creates a new subscription manager.
    20  // Args:
    21  // - networkUnderlay: the Underlay interface of the network layer.
    22  // Returns:
    23  // - a new subscription manager.
    24  func NewChannelSubscriptionManager(underlay network.Underlay) *ChannelSubscriptionManager {
    25  	return &ChannelSubscriptionManager{
    26  		engines:         make(map[channels.Channel]network.MessageProcessor),
    27  		networkUnderlay: underlay,
    28  	}
    29  }
    30  
    31  // Register registers an engine on the channel into the subscription manager.
    32  func (sm *ChannelSubscriptionManager) Register(channel channels.Channel, engine network.MessageProcessor) error {
    33  	sm.mu.Lock()
    34  	defer sm.mu.Unlock()
    35  
    36  	// channel should be registered only once.
    37  	_, ok := sm.engines[channel]
    38  	if ok {
    39  		return fmt.Errorf("subscriptionManager: channel already registered: %s", channel)
    40  	}
    41  
    42  	// registers the channel with the networkUnderlay to let networkUnderlay start receiving messages
    43  	// TODO: subscribe function should be replaced by a better abstraction of the network.
    44  	err := sm.networkUnderlay.Subscribe(channel)
    45  	if err != nil {
    46  		return fmt.Errorf("subscriptionManager: failed to subscribe to channel %s: %w", channel, err)
    47  	}
    48  
    49  	// saves the engine for the provided channel
    50  	sm.engines[channel] = engine
    51  
    52  	return nil
    53  }
    54  
    55  // Unregister removes the engine associated with a channel.
    56  func (sm *ChannelSubscriptionManager) Unregister(channel channels.Channel) error {
    57  	sm.mu.Lock()
    58  	defer sm.mu.Unlock()
    59  
    60  	// check if there is a registered engine for the given channel
    61  	_, ok := sm.engines[channel]
    62  	if !ok {
    63  		// if not found then there is nothing else to do
    64  		return nil
    65  	}
    66  
    67  	err := sm.networkUnderlay.Unsubscribe(channel)
    68  	if err != nil {
    69  		return fmt.Errorf("subscriptionManager: failed to unregister from channel %s: %w", channel, err)
    70  	}
    71  
    72  	delete(sm.engines, channel)
    73  
    74  	return nil
    75  }
    76  
    77  // GetEngine returns engine associated with a channel.
    78  func (sm *ChannelSubscriptionManager) GetEngine(channel channels.Channel) (network.MessageProcessor, error) {
    79  	sm.mu.RLock()
    80  	defer sm.mu.RUnlock()
    81  
    82  	eng, found := sm.engines[channel]
    83  	if !found {
    84  		return nil, fmt.Errorf("subscriptionManager: engine for channel %s not found", channel)
    85  	}
    86  	return eng, nil
    87  }
    88  
    89  // Channels returns all the channels registered in this subscription manager.
    90  func (sm *ChannelSubscriptionManager) Channels() channels.ChannelList {
    91  	sm.mu.RLock()
    92  	defer sm.mu.RUnlock()
    93  
    94  	channels := make(channels.ChannelList, 0)
    95  	for channel := range sm.engines {
    96  		channels = append(channels, channel)
    97  	}
    98  
    99  	return channels
   100  }