github.com/koko1123/flow-go-1@v0.29.6/network/p2p/subscription/subscriptionManager.go (about)

     1  package subscription
     2  
     3  import (
     4  	"fmt"
     5  	"sync"
     6  
     7  	"github.com/koko1123/flow-go-1/network"
     8  	"github.com/koko1123/flow-go-1/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  	mw      network.Middleware
    17  }
    18  
    19  func NewChannelSubscriptionManager(mw network.Middleware) *ChannelSubscriptionManager {
    20  	return &ChannelSubscriptionManager{
    21  		engines: make(map[channels.Channel]network.MessageProcessor),
    22  		mw:      mw,
    23  	}
    24  }
    25  
    26  // Register registers an engine on the channel into the subscription manager.
    27  func (sm *ChannelSubscriptionManager) Register(channel channels.Channel, engine network.MessageProcessor) error {
    28  	sm.mu.Lock()
    29  	defer sm.mu.Unlock()
    30  
    31  	// channel should be registered only once.
    32  	_, ok := sm.engines[channel]
    33  	if ok {
    34  		return fmt.Errorf("subscriptionManager: channel already registered: %s", channel)
    35  	}
    36  
    37  	// registers the channel with the middleware to let middleware start receiving messages
    38  	err := sm.mw.Subscribe(channel)
    39  	if err != nil {
    40  		return fmt.Errorf("subscriptionManager: failed to subscribe to channel %s: %w", channel, err)
    41  	}
    42  
    43  	// saves the engine for the provided channel
    44  	sm.engines[channel] = engine
    45  
    46  	return nil
    47  }
    48  
    49  // Unregister removes the engine associated with a channel.
    50  func (sm *ChannelSubscriptionManager) Unregister(channel channels.Channel) error {
    51  	sm.mu.Lock()
    52  	defer sm.mu.Unlock()
    53  
    54  	// check if there is a registered engine for the given channel
    55  	_, ok := sm.engines[channel]
    56  	if !ok {
    57  		// if not found then there is nothing else to do
    58  		return nil
    59  	}
    60  
    61  	err := sm.mw.Unsubscribe(channel)
    62  	if err != nil {
    63  		return fmt.Errorf("subscriptionManager: failed to unregister from channel %s", channel)
    64  	}
    65  
    66  	delete(sm.engines, channel)
    67  
    68  	return nil
    69  }
    70  
    71  // GetEngine returns engine associated with a channel.
    72  func (sm *ChannelSubscriptionManager) GetEngine(channel channels.Channel) (network.MessageProcessor, error) {
    73  	sm.mu.RLock()
    74  	defer sm.mu.RUnlock()
    75  
    76  	eng, found := sm.engines[channel]
    77  	if !found {
    78  		return nil, fmt.Errorf("subscriptionManager: engine for channel %s not found", channel)
    79  	}
    80  	return eng, nil
    81  }
    82  
    83  // Channels returns all the channels registered in this subscription manager.
    84  func (sm *ChannelSubscriptionManager) Channels() channels.ChannelList {
    85  	sm.mu.RLock()
    86  	defer sm.mu.RUnlock()
    87  
    88  	channels := make(channels.ChannelList, 0)
    89  	for channel := range sm.engines {
    90  		channels = append(channels, channel)
    91  	}
    92  
    93  	return channels
    94  }