github.com/iqoqo/nomad@v0.11.3-0.20200911112621-d7021c74d101/client/allocwatcher/group_alloc_watcher.go (about)

     1  package allocwatcher
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  
     7  	multierror "github.com/hashicorp/go-multierror"
     8  )
     9  
    10  type groupPrevAllocWatcher struct {
    11  	prevAllocs []PrevAllocWatcher
    12  	wg         sync.WaitGroup
    13  
    14  	// waiting and migrating are true when alloc runner is waiting on the
    15  	// prevAllocWatcher. Writers must acquire the waitingLock and readers
    16  	// should use the helper methods IsWaiting and IsMigrating.
    17  	waiting     bool
    18  	waitingLock sync.RWMutex
    19  }
    20  
    21  func NewGroupAllocWatcher(watchers ...PrevAllocWatcher) PrevAllocWatcher {
    22  	return &groupPrevAllocWatcher{
    23  		prevAllocs: watchers,
    24  	}
    25  }
    26  
    27  // Wait on the previous allocs to become terminal, exit, or, return due to
    28  // context termination. Usage of the groupPrevAllocWatcher requires that all
    29  // sub-watchers correctly handle context cancellation.
    30  // We may need to adjust this to use channels rather than a wait group, if we
    31  // wish to more strictly enforce timeouts.
    32  func (g *groupPrevAllocWatcher) Wait(ctx context.Context) error {
    33  	g.waitingLock.Lock()
    34  	g.waiting = true
    35  	g.waitingLock.Unlock()
    36  	defer func() {
    37  		g.waitingLock.Lock()
    38  		g.waiting = false
    39  		g.waitingLock.Unlock()
    40  	}()
    41  
    42  	var merr multierror.Error
    43  	var errmu sync.Mutex
    44  
    45  	g.wg.Add(len(g.prevAllocs))
    46  
    47  	for _, alloc := range g.prevAllocs {
    48  		go func(ctx context.Context, alloc PrevAllocWatcher) {
    49  			defer g.wg.Done()
    50  			err := alloc.Wait(ctx)
    51  			if err != nil {
    52  				errmu.Lock()
    53  				merr.Errors = append(merr.Errors, err)
    54  				errmu.Unlock()
    55  			}
    56  		}(ctx, alloc)
    57  	}
    58  
    59  	g.wg.Wait()
    60  
    61  	// Check ctx.Err first, to avoid returning an mErr of ctx.Err from prevAlloc
    62  	// Wait routines.
    63  	if err := ctx.Err(); err != nil {
    64  		return err
    65  	}
    66  
    67  	return merr.ErrorOrNil()
    68  }
    69  
    70  func (g *groupPrevAllocWatcher) IsWaiting() bool {
    71  	g.waitingLock.RLock()
    72  	defer g.waitingLock.RUnlock()
    73  
    74  	return g.waiting
    75  }