github.com/aergoio/aergo@v1.3.1/p2p/p2putil/channelpipe.go (about)

     1  package p2putil
     2  
     3  // ChannelPipe serve non blocking limited size channel.
     4  // It preserve input ordering, and not block caller unless it is Closed()
     5  // Tt must be called Open before using it, and Close for dispose resource.
     6  type ChannelPipe interface {
     7  	// Put item to pipe. it should be used after Open() method is called.
     8  	// It always returns true and guaranty that item is queued.
     9  	Put(item interface{}) bool
    10  	Out() <-chan interface{}
    11  	// Done should be called after get item from out channel
    12  	Done()
    13  
    14  	Open()
    15  	Close()
    16  }
    17  
    18  type channelPipe struct {
    19  	in   chan interface{}
    20  	out  chan interface{}
    21  	done chan interface{}
    22  
    23  	queue *PressableQueue
    24  	stop  chan interface{}
    25  
    26  	listener PipeEventListener
    27  }
    28  
    29  // PipeEventListener listen event of ChannelPipe
    30  type PipeEventListener interface {
    31  	// OnIn is called when item is queued
    32  	OnIn(element interface{})
    33  	// OnDrop is called when queued item is dropped and not out to channel receiver
    34  	OnDrop(element interface{})
    35  	// OnOut is called when queued item went to out channel (and will be sent to receiver)
    36  	OnOut(element interface{})
    37  }
    38  
    39  // NewDefaultChannelPipe create pipe to output channel out
    40  func NewDefaultChannelPipe(bufSize int, listener PipeEventListener) ChannelPipe {
    41  	return newDefaultChannelPipe(bufSize, listener)
    42  }
    43  
    44  // newDefaultChannelPipe create pipe to output channel out
    45  func newDefaultChannelPipe(bufSize int, listener PipeEventListener) *channelPipe {
    46  	if listener == nil {
    47  		listener = &StatListener{}
    48  	}
    49  	c := &channelPipe{
    50  		in:   make(chan interface{}),
    51  		out:  make(chan interface{}, 1),
    52  		done: make(chan interface{}),
    53  
    54  		queue: NewPressableQueue(bufSize),
    55  		stop:  make(chan interface{}),
    56  
    57  		listener: listener,
    58  	}
    59  
    60  	return c
    61  }
    62  
    63  func (c *channelPipe) Put(item interface{}) bool {
    64  	c.in <- item
    65  	return true
    66  }
    67  
    68  func (c *channelPipe) Out() <-chan interface{} {
    69  	return c.out
    70  }
    71  
    72  func (c *channelPipe) Done() {
    73  	c.done <- struct{}{}
    74  }
    75  
    76  func (c *channelPipe) Open() {
    77  	go c.run()
    78  }
    79  func (c *channelPipe) Close() {
    80  	c.stop <- struct{}{}
    81  }
    82  
    83  func (c *channelPipe) run() {
    84  LOOP:
    85  	for {
    86  		select {
    87  		case mo := <-c.in:
    88  			c.listener.OnIn(mo)
    89  			if len(c.out) == 0 {
    90  				if c.queue.Empty() {
    91  					c.pushToOut(mo)
    92  					continue LOOP
    93  				} else {
    94  					c.pushToOut(c.queue.Poll())
    95  					c.queue.Offer(mo)
    96  				}
    97  			} else {
    98  				if dropped := c.queue.Press(mo); dropped != nil {
    99  					c.listener.OnDrop(dropped)
   100  				}
   101  			}
   102  		case <-c.done:
   103  			// Next done will come, if out is not empty
   104  			if len(c.out) == 0 {
   105  				first := c.queue.Poll()
   106  				if first != nil {
   107  					c.pushToOut(first)
   108  				}
   109  			}
   110  		case <-c.stop:
   111  			break LOOP
   112  		}
   113  	}
   114  }
   115  
   116  func (c *channelPipe) pushToOut(e interface{}) {
   117  	c.out <- e
   118  	c.listener.OnOut(e)
   119  }