github.com/go-maxhub/gremlins@v1.0.1-0.20231227222204-b03a6a1e3e09/core/engine/workerpool/workerpool.go (about)

     1  /*
     2   * Copyright 2022 The Gremlins Authors
     3   *
     4   *    Licensed under the Apache License, Version 2.0 (the "License");
     5   *    you may not use this file except in compliance with the License.
     6   *    You may obtain a copy of the License at
     7   *
     8   *        http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   *    Unless required by applicable law or agreed to in writing, software
    11   *    distributed under the License is distributed on an "AS IS" BASIS,
    12   *    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   *    See the License for the specific language governing permissions and
    14   *    limitations under the License.
    15   */
    16  
    17  package workerpool
    18  
    19  import (
    20  	"runtime"
    21  	"sync"
    22  
    23  	"github.com/go-maxhub/gremlins/core/configuration"
    24  )
    25  
    26  // Executor is the unit of work that executes a task.
    27  type Executor interface {
    28  	Start(worker *Worker)
    29  }
    30  
    31  // Worker takes an executor and starts the actual executor.
    32  type Worker struct {
    33  	stopCh chan struct{}
    34  	Name   string
    35  	ID     int
    36  }
    37  
    38  // NewWorker instantiates a new worker with an ID and name.
    39  func NewWorker(id int, name string) *Worker {
    40  	return &Worker{
    41  		Name: name,
    42  		ID:   id,
    43  	}
    44  }
    45  
    46  // Start gets an executor queue and starts working on it.
    47  func (w *Worker) Start(executorQueue <-chan Executor) {
    48  	w.stopCh = make(chan struct{})
    49  	go func() {
    50  		for {
    51  			executor, ok := <-executorQueue
    52  			if !ok {
    53  				w.stopCh <- struct{}{}
    54  
    55  				break
    56  			}
    57  			executor.Start(w)
    58  		}
    59  	}()
    60  }
    61  
    62  func (w *Worker) stop() {
    63  	<-w.stopCh
    64  }
    65  
    66  // Pool manages and limits the number of concurrent Worker.
    67  type Pool struct {
    68  	queue   chan Executor
    69  	name    string
    70  	workers []*Worker
    71  	size    int
    72  }
    73  
    74  // Initialize creates a new Pool with a name and the number of parallel
    75  // workers it will use.
    76  func Initialize(name string) *Pool {
    77  	wNum := configuration.Get[int](configuration.UnleashWorkersKey)
    78  	intMode := configuration.Get[bool](configuration.UnleashIntegrationMode)
    79  
    80  	p := &Pool{
    81  		size: size(wNum, intMode),
    82  		name: name,
    83  	}
    84  	p.workers = []*Worker{}
    85  	for i := 0; i < p.size; i++ {
    86  		w := NewWorker(i, p.name)
    87  		p.workers = append(p.workers, w)
    88  	}
    89  	p.queue = make(chan Executor, 1)
    90  
    91  	return p
    92  }
    93  
    94  func size(wNum int, intMode bool) int {
    95  	if wNum == 0 {
    96  		wNum = runtime.NumCPU()
    97  	}
    98  	if intMode && wNum > 1 {
    99  		wNum /= 2
   100  	}
   101  
   102  	return wNum
   103  }
   104  
   105  // AppendExecutor adds a new Executor to the queue of Executor to be processed.
   106  func (p *Pool) AppendExecutor(executor Executor) {
   107  	p.queue <- executor
   108  }
   109  
   110  // Start the Pool.
   111  func (p *Pool) Start() {
   112  	for _, w := range p.workers {
   113  		w.Start(p.queue)
   114  	}
   115  }
   116  
   117  // Stop the Pool and wait for all the pending Worker to complete.
   118  func (p *Pool) Stop() {
   119  	close(p.queue)
   120  	var wg sync.WaitGroup
   121  	for _, worker := range p.workers {
   122  		wg.Add(1)
   123  		go func(w *Worker) {
   124  			defer wg.Done()
   125  			w.stop()
   126  		}(worker)
   127  	}
   128  	wg.Wait()
   129  }
   130  
   131  // ActiveWorkers gives the number of active workers on the Pool.
   132  func (p *Pool) ActiveWorkers() int {
   133  	return len(p.workers)
   134  }