github.com/searKing/golang/go@v1.2.117/sync/walk.go (about)

     1  // Copyright 2021 The searKing Author. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package sync
     6  
     7  import (
     8  	"context"
     9  	"sync"
    10  
    11  	"github.com/searKing/golang/go/time/rate"
    12  )
    13  
    14  type Walk struct {
    15  	Burst int // Burst will be set to 1 if less than 1
    16  
    17  	wg sync.WaitGroup
    18  
    19  	walkError     error
    20  	walkErrorOnce sync.Once
    21  	walkErrorMu   sync.Mutex
    22  }
    23  
    24  // WalkFunc is the type of the function called for each task processed
    25  // by Walk. The path argument contains the argument to Walk as a task.
    26  //
    27  // In the case of an error, the info argument will be nil. If an error
    28  // is returned, processing stops.
    29  type WalkFunc func(task any) error
    30  
    31  // Walk will consume all tasks parallel and block until ctx.Done() or taskChan is closed.
    32  // Walk returns a channel that's closed when work done on behalf of this
    33  // walk should be canceled. Done may return nil if this walk can
    34  // never be canceled. Successive calls to Done return the same value.
    35  // The close of the Done channel may happen asynchronously,
    36  // after the cancel function returns.
    37  func (p *Walk) Walk(ctx context.Context, taskChan <-chan any, procFn WalkFunc) (doneC <-chan struct{}) {
    38  	done := make(chan struct{})
    39  	if p.Burst <= 0 {
    40  		p.Burst = 1
    41  	}
    42  	p.wg.Add(1)
    43  
    44  	go func() {
    45  		defer close(done)
    46  		defer p.wg.Done()
    47  		ctx, cancel := context.WithCancel(ctx)
    48  		defer cancel()
    49  		limiter := rate.NewFullBurstLimiter(p.Burst)
    50  		for {
    51  			// Wait blocks until lim permits 1 events to happen.
    52  			if err := limiter.Wait(ctx); err != nil {
    53  				p.TrySetError(err)
    54  				return
    55  			}
    56  
    57  			select {
    58  			case task, ok := <-taskChan:
    59  				if !ok {
    60  					return
    61  				}
    62  				p.wg.Add(1)
    63  				go func() {
    64  					defer limiter.PutToken()
    65  					defer p.wg.Done()
    66  					if err := procFn(task); err != nil {
    67  						p.TrySetError(err)
    68  						return
    69  					}
    70  				}()
    71  			case <-ctx.Done():
    72  				p.TrySetError(ctx.Err())
    73  				return
    74  			}
    75  		}
    76  	}()
    77  	return done
    78  }
    79  
    80  // Wait blocks until the WaitGroup counter is zero.
    81  func (p *Walk) Wait() error {
    82  	p.wg.Wait()
    83  	return p.Error()
    84  }
    85  
    86  func (p *Walk) Error() error {
    87  	p.walkErrorMu.Lock()
    88  	defer p.walkErrorMu.Unlock()
    89  	return p.walkError
    90  }
    91  func (p *Walk) TrySetError(err error) {
    92  	p.walkErrorOnce.Do(func() {
    93  		p.walkErrorMu.Lock()
    94  		defer p.walkErrorMu.Unlock()
    95  		p.walkError = err
    96  	})
    97  }