github.com/anycable/anycable-go@v1.5.1/utils/gopool.go (about)

     1  package utils
     2  
     3  import (
     4  	"fmt"
     5  	"time"
     6  )
     7  
     8  // ErrScheduleTimeout returned by Pool to indicate that there no free
     9  // goroutines during some period of time.
    10  var ErrScheduleTimeout = fmt.Errorf("schedule error: timed out")
    11  
    12  // How many tasks should a worker perform before "re-starting" the goroutine
    13  // See https://adtac.in/2021/04/23/note-on-worker-pools-in-go.html
    14  const workerRespawnThreshold = 1 << 16
    15  
    16  // GoPool contains logic of goroutine reuse.
    17  // Copied from https://github.com/gobwas/ws-examples/blob/master/src/gopool/pool.go
    18  type GoPool struct {
    19  	name string
    20  	size int
    21  	sem  chan struct{}
    22  	work chan func()
    23  }
    24  
    25  var initializedPools []*GoPool = make([]*GoPool, 0)
    26  
    27  // Return all active pools
    28  func AllPools() []*GoPool {
    29  	return initializedPools
    30  }
    31  
    32  // NewGoPool creates new goroutine pool with given size.
    33  // Start size defaults to 20% of the max size but not greater than 1024.
    34  // Queue size defaults to 50% of the max size.
    35  func NewGoPool(name string, size int) *GoPool {
    36  	queue := min(size/2, 1)
    37  
    38  	spawn := min(max(size/5, 1), 1024)
    39  
    40  	p := &GoPool{
    41  		name: name,
    42  		size: size,
    43  		sem:  make(chan struct{}, size),
    44  		work: make(chan func(), queue),
    45  	}
    46  
    47  	for i := 0; i < spawn; i++ {
    48  		p.sem <- struct{}{}
    49  		go p.worker(func() {})
    50  	}
    51  
    52  	initializedPools = append(initializedPools, p)
    53  
    54  	return p
    55  }
    56  
    57  func (p *GoPool) Name() string {
    58  	return p.name
    59  }
    60  
    61  func (p *GoPool) Size() int {
    62  	return p.size
    63  }
    64  
    65  // Schedule schedules task to be executed over pool's workers.
    66  func (p *GoPool) Schedule(task func()) {
    67  	p.schedule(task, nil) // nolint:errcheck
    68  }
    69  
    70  // ScheduleTimeout schedules task to be executed over pool's workers.
    71  // It returns ErrScheduleTimeout when no free workers met during given timeout.
    72  func (p *GoPool) ScheduleTimeout(timeout time.Duration, task func()) error {
    73  	return p.schedule(task, time.After(timeout))
    74  }
    75  
    76  func (p *GoPool) schedule(task func(), timeout <-chan time.Time) error {
    77  	select {
    78  	case <-timeout:
    79  		return ErrScheduleTimeout
    80  	case p.work <- task:
    81  		return nil
    82  	case p.sem <- struct{}{}:
    83  		go p.worker(task)
    84  		return nil
    85  	}
    86  }
    87  
    88  func (p *GoPool) worker(task func()) {
    89  	counter := 1
    90  	defer func() { <-p.sem }()
    91  
    92  	task()
    93  
    94  	for task := range p.work {
    95  		task()
    96  		counter++
    97  
    98  		if counter >= workerRespawnThreshold {
    99  			return
   100  		}
   101  	}
   102  }