github.com/kubeshop/testkube@v1.17.23/pkg/workerpool/service.go (about)

     1  package workerpool
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  
     7  	"github.com/kubeshop/testkube/pkg/api/v1/testkube"
     8  )
     9  
    10  // Runnable is an interface of runnable objects
    11  type Runnable interface {
    12  	testkube.Test | testkube.TestSuite
    13  }
    14  
    15  // Requestable is an interface of requestable objects
    16  type Requestable interface {
    17  	testkube.ExecutionRequest | testkube.TestSuiteExecutionRequest
    18  }
    19  
    20  // Returnable is an interface of returnable objects
    21  type Returnable interface {
    22  	testkube.Execution | testkube.TestSuiteExecution
    23  }
    24  
    25  // ExecuteFn is a function type for executing runnable and requestable parameters with returnable results
    26  type ExecuteFn[R Runnable, T Requestable, E Returnable] func(ctx context.Context, object R, options T) (result E, err error)
    27  
    28  // Request contains request parameters and invocation method
    29  type Request[R Runnable, T Requestable, E Returnable] struct {
    30  	Object  R
    31  	Options T
    32  	ExecFn  ExecuteFn[R, T, E]
    33  }
    34  
    35  // Response contains result details
    36  type Response[E Returnable] struct {
    37  	Result E
    38  	Err    error
    39  }
    40  
    41  // execute is a method wrapper for ExecFn execution
    42  func (r Request[R, T, E]) execute(ctx context.Context) Response[E] {
    43  	result, err := r.ExecFn(ctx, r.Object, r.Options)
    44  	if err != nil {
    45  		return Response[E]{
    46  			Err: err,
    47  		}
    48  	}
    49  
    50  	return Response[E]{
    51  		Result: result,
    52  	}
    53  }
    54  
    55  // Service is a worker pool service
    56  type Service[R Runnable, T Requestable, E Returnable] struct {
    57  	concurrencyLevel int
    58  	requests         chan Request[R, T, E]
    59  	responses        chan Response[E]
    60  }
    61  
    62  // New is a constructor for worker pool service
    63  func New[R Runnable, T Requestable, E Returnable](concurrencyLevel int) Service[R, T, E] {
    64  	return Service[R, T, E]{
    65  		concurrencyLevel: concurrencyLevel,
    66  		requests:         make(chan Request[R, T, E], concurrencyLevel),
    67  		responses:        make(chan Response[E], concurrencyLevel),
    68  	}
    69  }
    70  
    71  // Run is a method to run worker pool
    72  func (s Service[R, T, E]) Run(ctx context.Context) {
    73  	var wg sync.WaitGroup
    74  
    75  	for i := 0; i < s.concurrencyLevel; i++ {
    76  		wg.Add(1)
    77  		go worker(ctx, &wg, s.requests, s.responses)
    78  	}
    79  
    80  	wg.Wait()
    81  	close(s.responses)
    82  }
    83  
    84  // GetResponse return reponses of method execution
    85  func (s Service[R, T, E]) GetResponses() <-chan Response[E] {
    86  	return s.responses
    87  }
    88  
    89  // SendRequests sends requests to workers
    90  func (s Service[R, T, E]) SendRequests(requests []Request[R, T, E]) {
    91  	for i := range requests {
    92  		s.requests <- requests[i]
    93  	}
    94  	close(s.requests)
    95  }
    96  
    97  // worker is a worker pool method
    98  func worker[R Runnable, T Requestable, E Returnable](ctx context.Context, wg *sync.WaitGroup,
    99  	requests <-chan Request[R, T, E], responses chan<- Response[E]) {
   100  	defer wg.Done()
   101  	for {
   102  		select {
   103  		case request, ok := <-requests:
   104  			if !ok {
   105  				return
   106  			}
   107  
   108  			responses <- request.execute(ctx)
   109  		case <-ctx.Done():
   110  			return
   111  		}
   112  	}
   113  }