github.com/mailgun/holster/v4@v4.20.0/syncutil/waitgroup.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  type WaitGroup struct {
    21  	wg    sync.WaitGroup
    22  	mutex sync.Mutex
    23  	errs  []error
    24  	done  chan struct{}
    25  }
    26  
    27  // Run a routine and collect errors if any
    28  func (wg *WaitGroup) Run(callBack func(interface{}) error, data interface{}) {
    29  	wg.wg.Add(1)
    30  	go func() {
    31  		err := callBack(data)
    32  		if err == nil {
    33  			wg.wg.Done()
    34  			return
    35  		}
    36  		wg.mutex.Lock()
    37  		wg.errs = append(wg.errs, err)
    38  		wg.wg.Done()
    39  		wg.mutex.Unlock()
    40  	}()
    41  }
    42  
    43  // Execute a long running routine
    44  func (wg *WaitGroup) Go(cb func()) {
    45  	wg.wg.Add(1)
    46  	go func() {
    47  		cb()
    48  		wg.wg.Done()
    49  	}()
    50  }
    51  
    52  // Run a goroutine in a loop continuously, if the callBack returns false the loop is broken.
    53  // `Until()` differs from `Loop()` in that if the `Stop()` is called on the WaitGroup
    54  // the `done` channel is closed. Implementations of the callBack function can listen
    55  // for the close to indicate a stop was requested.
    56  func (wg *WaitGroup) Until(callBack func(done chan struct{}) bool) {
    57  	wg.mutex.Lock()
    58  	if wg.done == nil {
    59  		wg.done = make(chan struct{})
    60  	}
    61  	wg.mutex.Unlock()
    62  
    63  	wg.wg.Add(1)
    64  	go func() {
    65  		for {
    66  			if !callBack(wg.done) {
    67  				wg.wg.Done()
    68  				break
    69  			}
    70  		}
    71  	}()
    72  }
    73  
    74  // Stop closes the done channel passed into `Until()` calls and waits for
    75  // the `Until()` callBack to return false.
    76  func (wg *WaitGroup) Stop() {
    77  	wg.mutex.Lock()
    78  	defer wg.mutex.Unlock()
    79  
    80  	if wg.done != nil {
    81  		close(wg.done)
    82  	}
    83  	wg.wg.Wait()
    84  	wg.done = nil
    85  }
    86  
    87  // Run a goroutine in a loop continuously, if the callBack returns false the loop is broken
    88  func (wg *WaitGroup) Loop(callBack func() bool) {
    89  	wg.wg.Add(1)
    90  	go func() {
    91  		for {
    92  			if !callBack() {
    93  				wg.wg.Done()
    94  				break
    95  			}
    96  		}
    97  	}()
    98  }
    99  
   100  // Wait for all the routines to complete and return any errors collected
   101  func (wg *WaitGroup) Wait() []error {
   102  	wg.wg.Wait()
   103  
   104  	wg.mutex.Lock()
   105  	defer wg.mutex.Unlock()
   106  
   107  	if len(wg.errs) == 0 {
   108  		return nil
   109  	}
   110  	return wg.errs
   111  }