github.com/jwowillo/pipe@v1.2.0/pipe.go (about)

     1  // Package pipe allows Stages of functions to easily be assembled into concurent
     2  // pipes.
     3  package pipe
     4  
     5  import (
     6  	"sync"
     7  )
     8  
     9  // Item to be handled by a Stage.
    10  type Item interface{}
    11  
    12  // Stage handles an Item.
    13  type Stage interface {
    14  	Handle(Item) Item
    15  }
    16  
    17  // StageFunc converts a function to a Stage.
    18  type StageFunc func(Item) Item
    19  
    20  // Handle applies the function to the Item.
    21  func (f StageFunc) Handle(x Item) Item {
    22  	return f(x)
    23  }
    24  
    25  // Pipe connects Stages so many Items can be processed by the Stages in order
    26  // concurrently.
    27  type Pipe struct {
    28  	m      sync.Mutex
    29  	count  int
    30  	stages []Stage
    31  	links  []chan Item
    32  }
    33  
    34  // New Pipe with all the Stages connected in the order given.
    35  func New(ss ...Stage) *Pipe {
    36  	return &Pipe{
    37  		m:      sync.Mutex{},
    38  		stages: append([]Stage{}, ss...),
    39  		links:  make([]chan Item, len(ss)+1),
    40  	}
    41  }
    42  
    43  // Receive the Item into the beginning of the Pipe.
    44  func (p *Pipe) Receive(x Item) {
    45  	p.m.Lock()
    46  	defer p.m.Unlock()
    47  	if p.isEmpty() {
    48  		p.start()
    49  	}
    50  	p.count++
    51  	go func() { p.links[0] <- x }()
    52  }
    53  
    54  // Deliver the item from the end of the Pipe once it's ready.
    55  func (p *Pipe) Deliver() Item {
    56  	p.m.Lock()
    57  	defer p.m.Unlock()
    58  	x := <-p.links[len(p.links)-1]
    59  	p.count--
    60  	if p.isEmpty() {
    61  		p.stop()
    62  	}
    63  	return x
    64  }
    65  
    66  func (p *Pipe) isEmpty() bool {
    67  	return p.count == 0
    68  }
    69  
    70  func (p *Pipe) start() {
    71  	p.links[0] = make(chan Item)
    72  	for i, stage := range p.stages {
    73  		p.links[i+1] = make(chan Item)
    74  		go func(receive <-chan Item, send chan<- Item, stage Stage) {
    75  			for x := range receive {
    76  				go func(x Item) {
    77  					send <- stage.Handle(x)
    78  				}(x)
    79  			}
    80  		}(p.links[i], p.links[i+1], stage)
    81  	}
    82  }
    83  
    84  func (p *Pipe) stop() {
    85  	for _, link := range p.links {
    86  		close(link)
    87  	}
    88  }
    89  
    90  // Process all the Items with the Pipe.
    91  func Process(p *Pipe, xs ...Item) []Item {
    92  	for _, x := range xs {
    93  		p.Receive(x)
    94  	}
    95  	out := make([]Item, len(xs))
    96  	for i := range out {
    97  		out[i] = p.Deliver()
    98  	}
    99  	return out
   100  }