github.com/defanghe/fabric@v2.1.1+incompatible/gossip/comm/demux.go (about)

     1  /*
     2  Copyright IBM Corp. All Rights Reserved.
     3  
     4  SPDX-License-Identifier: Apache-2.0
     5  */
     6  
     7  package comm
     8  
     9  import (
    10  	"sync"
    11  
    12  	"github.com/hyperledger/fabric/gossip/common"
    13  )
    14  
    15  // ChannelDeMultiplexer is a struct that can receive channel registrations (AddChannel)
    16  // and publications (DeMultiplex) and it broadcasts the publications to registrations
    17  // according to their predicate. Can only be closed once and never open after a close.
    18  type ChannelDeMultiplexer struct {
    19  	// lock protects everything below it.
    20  	lock   sync.Mutex
    21  	closed bool // one way boolean from false -> true
    22  	stopCh chan struct{}
    23  	// deMuxInProgress keeps track of any calls to DeMultiplex
    24  	// that are still being handled. This is used to determine
    25  	// when it is safe to close all of the tracked channels
    26  	deMuxInProgress sync.WaitGroup
    27  	channels        []*channel
    28  }
    29  
    30  // NewChannelDemultiplexer creates a new ChannelDeMultiplexer
    31  func NewChannelDemultiplexer() *ChannelDeMultiplexer {
    32  	return &ChannelDeMultiplexer{stopCh: make(chan struct{})}
    33  }
    34  
    35  type channel struct {
    36  	pred common.MessageAcceptor
    37  	ch   chan<- interface{}
    38  }
    39  
    40  // Close closes this channel, which makes all channels registered before
    41  // to close as well.
    42  func (m *ChannelDeMultiplexer) Close() {
    43  	m.lock.Lock()
    44  	if m.closed {
    45  		m.lock.Unlock()
    46  		return
    47  	}
    48  	m.closed = true
    49  	close(m.stopCh)
    50  	m.deMuxInProgress.Wait()
    51  	for _, ch := range m.channels {
    52  		close(ch.ch)
    53  	}
    54  	m.channels = nil
    55  	m.lock.Unlock()
    56  }
    57  
    58  // AddChannel registers a channel with a certain predicate. AddChannel
    59  // returns a read-only channel that will produce values that are
    60  // matched by the predicate function.
    61  //
    62  // If the DeMultiplexer is closed, the channel returned will be closed
    63  // to prevent users of the channel from waiting on the channel.
    64  func (m *ChannelDeMultiplexer) AddChannel(predicate common.MessageAcceptor) <-chan interface{} {
    65  	m.lock.Lock()
    66  	if m.closed { // closed once, can't put anything more in.
    67  		m.lock.Unlock()
    68  		ch := make(chan interface{})
    69  		close(ch)
    70  		return ch
    71  	}
    72  	bidirectionalCh := make(chan interface{}, 10)
    73  	// Assignment to channel converts bidirectionalCh to send-only.
    74  	// Return converts bidirectionalCh to a receive-only.
    75  	ch := &channel{ch: bidirectionalCh, pred: predicate}
    76  	m.channels = append(m.channels, ch)
    77  	m.lock.Unlock()
    78  	return bidirectionalCh
    79  }
    80  
    81  // DeMultiplex broadcasts the message to all channels that were returned
    82  // by AddChannel calls and that hold the respected predicates.
    83  //
    84  // Blocks if any one channel that would receive msg has a full buffer.
    85  func (m *ChannelDeMultiplexer) DeMultiplex(msg interface{}) {
    86  	m.lock.Lock()
    87  	if m.closed {
    88  		m.lock.Unlock()
    89  		return
    90  	}
    91  	channels := m.channels
    92  	m.deMuxInProgress.Add(1)
    93  	m.lock.Unlock()
    94  
    95  	for _, ch := range channels {
    96  		if ch.pred(msg) {
    97  			select {
    98  			case <-m.stopCh:
    99  				m.deMuxInProgress.Done()
   100  				return // stopping
   101  			case ch.ch <- msg:
   102  			}
   103  		}
   104  	}
   105  	m.deMuxInProgress.Done()
   106  }