github.com/juju/juju@v0.0.0-20240430160146-1752b71fcf00/core/workerpool/workerpool_test.go (about)

     1  // Copyright 2021 Canonical Ltd.
     2  // Licensed under the AGPLv3, see LICENCE file for details.
     3  
     4  package workerpool
     5  
     6  import (
     7  	"sort"
     8  	"strings"
     9  	"sync"
    10  	"time"
    11  
    12  	"github.com/juju/errors"
    13  	"github.com/juju/loggo"
    14  	jc "github.com/juju/testing/checkers"
    15  	gc "gopkg.in/check.v1"
    16  
    17  	coretesting "github.com/juju/juju/testing"
    18  )
    19  
    20  var _ = gc.Suite(&ProvisionerWorkerPoolSuite{})
    21  
    22  type ProvisionerWorkerPoolSuite struct {
    23  }
    24  
    25  func (s *ProvisionerWorkerPoolSuite) TestProcessMoreTasksThanWorkers(c *gc.C) {
    26  	doneCh := make(chan struct{}, 10)
    27  	wp := NewWorkerPool(loggo.GetLogger("test"), 5)
    28  	c.Assert(wp.Size(), gc.Equals, 5)
    29  
    30  	for i := 0; i < 10; i++ {
    31  		task := Task{
    32  			Type: "alien invasion",
    33  			Process: func() error {
    34  				doneCh <- struct{}{}
    35  				return nil
    36  			},
    37  		}
    38  
    39  		select {
    40  		case wp.Queue() <- task:
    41  		case <-time.After(coretesting.LongWait):
    42  			c.Fatal("timeout waiting to enqueue task")
    43  		}
    44  	}
    45  
    46  	for i := 0; i < 10; i++ {
    47  		select {
    48  		case <-doneCh: // task ACK'd
    49  		case <-time.After(coretesting.LongWait):
    50  			c.Fatalf("timeout waiting for task %d to complete", i)
    51  		}
    52  	}
    53  
    54  	// Shutdown the pool and ensure that no errors got reported.
    55  	c.Assert(wp.Close(), jc.ErrorIsNil)
    56  }
    57  
    58  func (s *ProvisionerWorkerPoolSuite) TestConsolidateErrors(c *gc.C) {
    59  	var (
    60  		wg      sync.WaitGroup
    61  		barrier = make(chan struct{})
    62  		wp      = NewWorkerPool(loggo.GetLogger("test"), 3)
    63  	)
    64  
    65  	wg.Add(3)
    66  	for i := 0; i < 3; i++ {
    67  		// The even-numbered workers emit an error while the odd ones
    68  		// do not.
    69  		var expErr error
    70  		var task = Task{Type: "alien abduction"}
    71  		if i%2 == 0 {
    72  			expErr = errors.Errorf("worker %d out of fuel error", i)
    73  		}
    74  		task.Process = func() error {
    75  			// Signal that we are ready to process and wait for
    76  			// barrier to be released.
    77  			wg.Done()
    78  			<-barrier
    79  
    80  			return expErr
    81  		}
    82  
    83  		select {
    84  		case wp.Queue() <- task:
    85  		case <-time.After(coretesting.LongWait):
    86  			c.Fatal("timeout waiting to enqueue task")
    87  		}
    88  	}
    89  
    90  	// Wait for each worker to grab a taskuest and the lift the barrier
    91  	wg.Wait()
    92  	close(barrier)
    93  
    94  	// Wait for all workers to complete and exit
    95  	err := wp.Close()
    96  	c.Assert(err, gc.Not(gc.IsNil), gc.Commentf("expected individual worker errors to be consolidated into a single error"))
    97  
    98  	// Errors can appear in any order so we need to sort them first
    99  	errLines := strings.Split(err.Error(), "\n")
   100  	c.Assert(len(errLines), gc.Equals, 2, gc.Commentf("expected 2 errors to be consolidated"))
   101  	sort.Strings(errLines)
   102  	c.Assert(errLines[0], gc.Equals, "worker 0 out of fuel error")
   103  	c.Assert(errLines[1], gc.Equals, "worker 2 out of fuel error")
   104  }