github.com/mailgun/holster/v4@v4.20.0/syncutil/fanout.go (about)

     1  /*
     2  Copyright 2017 Mailgun Technologies Inc
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8  	http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  package syncutil
    17  
    18  import "sync"
    19  
    20  // FanOut spawns a new go-routine each time `Run()` is called until `size` is reached,
    21  // subsequent calls to `Run()` will block until previously `Run()` routines have completed.
    22  // Allowing the user to control how many routines will run simultaneously. `Wait()` then
    23  // collects any errors from the routines once they have all completed.
    24  type FanOut struct {
    25  	errChan chan error
    26  	size    chan bool
    27  	errs    []error
    28  	wg      sync.WaitGroup
    29  }
    30  
    31  func NewFanOut(size int) *FanOut {
    32  	// They probably want no concurrency
    33  	if size == 0 {
    34  		size = 1
    35  	}
    36  
    37  	pool := FanOut{
    38  		errChan: make(chan error, size),
    39  		size:    make(chan bool, size),
    40  		errs:    make([]error, 0),
    41  	}
    42  	pool.start()
    43  	return &pool
    44  }
    45  
    46  func (p *FanOut) start() {
    47  	p.wg.Add(1)
    48  	go func() {
    49  		defer p.wg.Done()
    50  		for err := range p.errChan {
    51  			p.errs = append(p.errs, err)
    52  		}
    53  	}()
    54  }
    55  
    56  // Run a new routine with an optional data value
    57  func (p *FanOut) Run(callBack func(interface{}) error, data interface{}) {
    58  	p.size <- true
    59  	go func() {
    60  		err := callBack(data)
    61  		if err != nil {
    62  			p.errChan <- err
    63  		}
    64  		<-p.size
    65  	}()
    66  }
    67  
    68  // Wait for all the routines to complete and return any errors
    69  func (p *FanOut) Wait() []error {
    70  	// Wait for all the routines to complete
    71  	for i := 0; i < cap(p.size); i++ {
    72  		p.size <- true
    73  	}
    74  	// Close the err channel
    75  	if p.errChan != nil {
    76  		close(p.errChan)
    77  	}
    78  
    79  	// Wait until the error collector routine is complete
    80  	p.wg.Wait()
    81  
    82  	// If there are no errors
    83  	if len(p.errs) == 0 {
    84  		return nil
    85  	}
    86  	return p.errs
    87  }