lab.nexedi.com/kirr/go123@v0.0.0-20240207185015-8299741fa871/xsync/xsync.go (about)

     1  // Copyright (C) 2019-2021  Nexedi SA and Contributors.
     2  //                          Kirill Smelkov <kirr@nexedi.com>
     3  //
     4  // This program is free software: you can Use, Study, Modify and Redistribute
     5  // it under the terms of the GNU General Public License version 3, or (at your
     6  // option) any later version, as published by the Free Software Foundation.
     7  //
     8  // You can also Link and Combine this program with other software covered by
     9  // the terms of any of the Free Software licenses or any of the Open Source
    10  // Initiative approved licenses and Convey the resulting work. Corresponding
    11  // source of such a combination shall include the source code for all other
    12  // software used.
    13  //
    14  // This program is distributed WITHOUT ANY WARRANTY; without even the implied
    15  // warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
    16  //
    17  // See COPYING file for full licensing terms.
    18  // See https://www.nexedi.com/licensing for rationale and options.
    19  
    20  // Package xsync complements standard package sync.
    21  //
    22  //   - `WorkGroup` allows to spawn group of goroutines working on a common task.
    23  //
    24  // Functionality provided by xsync package is also provided by Pygolang(*) in its
    25  // standard package sync.
    26  //
    27  // (*) https://pypi.org/project/pygolang
    28  package xsync
    29  
    30  import (
    31  	"context"
    32  	"sync"
    33  )
    34  
    35  // WorkGroup represents group of goroutines working on a common task.
    36  //
    37  // Use .Go() to spawn goroutines, and .Wait() to wait for all of them to
    38  // complete, for example:
    39  //
    40  //	wg := xsync.NewWorkGroup(ctx)
    41  //	wg.Go(f1)
    42  //	wg.Go(f2)
    43  //	err := wg.Wait()
    44  //
    45  // Every spawned function accepts context related to the whole work and derived
    46  // from ctx used to initialize WorkGroup, for example:
    47  //
    48  //	func f1(ctx context.Context) error {
    49  //	    ...
    50  //	}
    51  //
    52  // Whenever a function returns error, the work context is canceled indicating
    53  // to other spawned goroutines that they have to cancel their work. .Wait()
    54  // waits for all spawned goroutines to complete and returns error, if any, from
    55  // the first failed subtask.
    56  //
    57  // NOTE if spawned function panics, the panic is currently _not_ propagated to .Wait().
    58  //
    59  // WorkGroup is modelled after https://godoc.org/golang.org/x/sync/errgroup but
    60  // is not equal to it.
    61  type WorkGroup struct {
    62  	ctx    context.Context // workers are spawned under ctx
    63  	cancel func()          // aborts ctx
    64  	waitg  sync.WaitGroup  // wait group for workers
    65  	mu     sync.Mutex
    66  	err    error           // error of the first failed worker
    67  }
    68  
    69  // NewWorkGroup creates new WorkGroup working under ctx.
    70  //
    71  // See WorkGroup documentation for details.
    72  func NewWorkGroup(ctx context.Context) *WorkGroup {
    73  	g := &WorkGroup{}
    74  	g.ctx, g.cancel = context.WithCancel(ctx)
    75  	return g
    76  }
    77  
    78  // Go spawns new worker under workgroup.
    79  //
    80  // See WorkGroup documentation for details.
    81  func (g *WorkGroup) Go(f func(context.Context) error) {
    82  	g.waitg.Add(1)
    83  	go func() {
    84  		defer g.waitg.Done()
    85  
    86  		err := f(g.ctx)
    87  		if err == nil {
    88  			return
    89  		}
    90  
    91  		g.mu.Lock()
    92  		defer g.mu.Unlock()
    93  
    94  		if g.err == nil {
    95  			// this goroutine is the first failed task
    96  			g.err = err
    97  			g.cancel()
    98  		}
    99  	}()
   100  }
   101  
   102  // Wait waits for all spawned workers to complete.
   103  //
   104  // It returns the error, if any, from the first failed worker.
   105  // See WorkGroup documentation for details.
   106  func (g *WorkGroup) Wait() error {
   107  	g.waitg.Wait()
   108  	g.cancel()
   109  	return g.err
   110  }