github.com/koko1123/flow-go-1@v0.29.6/engine/common/splitter/network/network.go (about)

     1  package network
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  	"sync"
     7  
     8  	"github.com/ipfs/go-datastore"
     9  	"github.com/libp2p/go-libp2p/core/protocol"
    10  	"github.com/rs/zerolog"
    11  
    12  	splitterEngine "github.com/koko1123/flow-go-1/engine/common/splitter"
    13  	"github.com/koko1123/flow-go-1/module/component"
    14  	"github.com/koko1123/flow-go-1/module/irrecoverable"
    15  	"github.com/koko1123/flow-go-1/module/util"
    16  	"github.com/koko1123/flow-go-1/network"
    17  	"github.com/koko1123/flow-go-1/network/channels"
    18  )
    19  
    20  // Network is the splitter network. It is a wrapper around the default network implementation
    21  // and should be passed in to engine constructors that require a network to register with.
    22  // When an engine is registered with the splitter network, a splitter engine is created for
    23  // the given channel (if one doesn't already exist) and the engine is registered with that
    24  // splitter engine. As a result, multiple engines can register with the splitter network on
    25  // the same channel and will each receive all events on that channel.
    26  type Network struct {
    27  	net       network.Network
    28  	mu        sync.RWMutex
    29  	log       zerolog.Logger
    30  	splitters map[channels.Channel]*splitterEngine.Engine // stores splitters for each channel
    31  	conduits  map[channels.Channel]network.Conduit        // stores conduits for all registered channels
    32  	*component.ComponentManager
    33  }
    34  
    35  var _ network.Network = (*Network)(nil)
    36  
    37  // NewNetwork returns a new splitter network.
    38  func NewNetwork(
    39  	net network.Network,
    40  	log zerolog.Logger,
    41  ) *Network {
    42  	n := &Network{
    43  		net:       net,
    44  		splitters: make(map[channels.Channel]*splitterEngine.Engine),
    45  		conduits:  make(map[channels.Channel]network.Conduit),
    46  		log:       log,
    47  	}
    48  
    49  	n.ComponentManager = component.NewComponentManagerBuilder().
    50  		AddWorker(func(ctx irrecoverable.SignalerContext, ready component.ReadyFunc) {
    51  			err := util.WaitClosed(ctx, n.net.Ready())
    52  
    53  			if err != nil {
    54  				return
    55  			}
    56  
    57  			ready()
    58  
    59  			<-ctx.Done()
    60  		}).Build()
    61  
    62  	return n
    63  }
    64  
    65  func (n *Network) RegisterBlobService(channel channels.Channel, store datastore.Batching, opts ...network.BlobServiceOption) (network.BlobService, error) {
    66  	return n.net.RegisterBlobService(channel, store, opts...)
    67  }
    68  
    69  func (n *Network) RegisterPingService(pid protocol.ID, provider network.PingInfoProvider) (network.PingService, error) {
    70  	return n.net.RegisterPingService(pid, provider)
    71  }
    72  
    73  // Register will subscribe the given engine with the spitter on the given channel, and all registered
    74  // engines will be notified with incoming messages on the channel.
    75  // The returned Conduit can be used to send messages to engines on other nodes subscribed to the same channel
    76  func (n *Network) Register(channel channels.Channel, engine network.MessageProcessor) (network.Conduit, error) {
    77  	n.mu.Lock()
    78  	defer n.mu.Unlock()
    79  
    80  	splitter, splitterExists := n.splitters[channel]
    81  	conduit, conduitExists := n.conduits[channel]
    82  
    83  	if splitterExists != conduitExists {
    84  		return nil, errors.New("inconsistent state detected")
    85  	}
    86  
    87  	channelRegistered := splitterExists && conduitExists
    88  
    89  	if !channelRegistered {
    90  		// create new splitter for the channel
    91  		splitter = splitterEngine.New(
    92  			n.log,
    93  			channel,
    94  		)
    95  
    96  		n.splitters[channel] = splitter
    97  	}
    98  
    99  	// register engine with splitter
   100  	splitter.RegisterEngine(engine)
   101  
   102  	if !channelRegistered {
   103  		var err error
   104  		conduit, err = n.net.Register(channel, splitter)
   105  
   106  		if err != nil {
   107  			// undo previous steps
   108  			splitter.UnregisterEngine(engine)
   109  			delete(n.splitters, channel)
   110  
   111  			return nil, fmt.Errorf("failed to register splitter engine on channel %s: %w", channel, err)
   112  		}
   113  
   114  		n.conduits[channel] = conduit
   115  	}
   116  
   117  	return conduit, nil
   118  }