github.com/nyan233/littlerpc@v0.4.6-0.20230316182519-0c8d5c48abaf/internal/pool/task_pool.go (about)

     1  package pool
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"sync"
     7  	"sync/atomic"
     8  	"time"
     9  )
    10  
    11  const (
    12  	MaxTaskPoolSize = 1024 * 16
    13  )
    14  
    15  type RecoverFunc func(poolId int, err interface{})
    16  
    17  // DynamicTaskPool
    18  // v0.10 -> v0.36 实现了简单的任务池
    19  // v0.38 -> now 实现了可自动扩容的Goroutine池和可拓展的接口
    20  type DynamicTaskPool[Key Hash] struct {
    21  	// buf chan
    22  	tasks     chan func()
    23  	recoverFn RecoverFunc
    24  	// 用于取消所有goroutine
    25  	ctx      context.Context
    26  	cancelFn context.CancelFunc
    27  	// 统计关闭的goroutine数量
    28  	wg *sync.WaitGroup
    29  	// 现在活跃的goroutine数量
    30  	liveSize int32
    31  	// 池的大小，也即活跃的goroutine数量
    32  	minSize int32
    33  	// 最大的池大小
    34  	maxSize    int32
    35  	idleTicker *time.Ticker
    36  	// 关闭的标志
    37  	closed int32
    38  	// 执行成功的统计
    39  	_           [128 - 88]byte
    40  	execSuccess uint64
    41  	_           [128 - 8]byte
    42  	// 执行失败的统计
    43  	execFailed uint64
    44  }
    45  
    46  func NewTaskPool[Key Hash](bufSize, minSize, maxSize int32, rf RecoverFunc) TaskPool[Key] {
    47  	pool := new(DynamicTaskPool[Key])
    48  	if bufSize > MaxTaskPoolSize {
    49  		bufSize = MaxTaskPoolSize
    50  	}
    51  	pool.tasks = make(chan func(), bufSize)
    52  	pool.recoverFn = rf
    53  	pool.minSize = minSize
    54  	pool.maxSize = maxSize
    55  	pool.idleTicker = time.NewTicker(time.Second * 90)
    56  	pool.wg = &sync.WaitGroup{}
    57  	pool.wg.Add(int(minSize))
    58  	pool.ctx, pool.cancelFn = context.WithCancel(context.Background())
    59  	pool.start()
    60  	return pool
    61  }
    62  
    63  func (p *DynamicTaskPool[Key]) start() {
    64  	for i := 0; i < int(p.minSize); i++ {
    65  		gIndex := atomic.AddInt32(&p.liveSize, 1)
    66  		go func() {
    67  			exec[struct{}, any, Key](p, gIndex, p.ctx.Done(), nil)
    68  		}()
    69  	}
    70  }
    71  
    72  func exec[T any, T2 any, Key Hash](p *DynamicTaskPool[Key], gIndex int32, done <-chan T, done2 <-chan T2) {
    73  	defer p.wg.Done()
    74  	defer atomic.AddInt32(&p.liveSize, -1)
    75  	iFunc := func(fn func()) {
    76  		defer func() {
    77  			if err := recover(); err != nil {
    78  				atomic.AddUint64(&p.execFailed, 1)
    79  				p.recoverFn(int(gIndex), err)
    80  			} else {
    81  				atomic.AddUint64(&p.execSuccess, 1)
    82  			}
    83  		}()
    84  		fn()
    85  	}
    86  	for {
    87  		select {
    88  		case fn := <-p.tasks:
    89  			iFunc(fn)
    90  		case <-done:
    91  			return
    92  		case <-done2:
    93  			// select一个为nil的chan将始终阻塞, 所以done2在start()函数调用时必须是 == nil
    94  			return
    95  		}
    96  	}
    97  }
    98  
    99  func (p *DynamicTaskPool[Key]) Push(key Key, fn func()) error {
   100  	if atomic.LoadInt32(&p.closed) == 1 {
   101  		return errors.New("pool already closed")
   102  	}
   103  	select {
   104  	case p.tasks <- fn:
   105  		break
   106  	default:
   107  		// 阻塞表示buf已满, 需要扩容Goroutine
   108  		for {
   109  			oldLiveSize := atomic.LoadInt32(&p.liveSize)
   110  			if oldLiveSize >= p.maxSize {
   111  				// 已到达Goroutine上限则不扩容, 等待可用
   112  				p.tasks <- fn
   113  				return nil
   114  			}
   115  			if atomic.CompareAndSwapInt32(&p.liveSize, oldLiveSize, oldLiveSize+1) {
   116  				p.wg.Add(1)
   117  				go func() {
   118  					p.idleTicker.Reset(time.Second * 90)
   119  					exec[time.Time, struct{}](p, oldLiveSize+1, p.idleTicker.C, p.ctx.Done())
   120  				}()
   121  				p.tasks <- fn
   122  				return nil
   123  			} else {
   124  				// retry
   125  				continue
   126  			}
   127  		}
   128  	}
   129  	return nil
   130  }
   131  
   132  func (p *DynamicTaskPool[Key]) Stop() error {
   133  	if !atomic.CompareAndSwapInt32(&p.closed, 0, 1) {
   134  		return errors.New("pool already closed")
   135  	}
   136  	p.cancelFn()
   137  	p.wg.Wait()
   138  	return nil
   139  }
   140  
   141  func (p *DynamicTaskPool[Key]) LiveSize() int {
   142  	return int(atomic.LoadInt32(&p.liveSize))
   143  }
   144  
   145  func (p *DynamicTaskPool[Key]) BufSize() int {
   146  	return len(p.tasks)
   147  }
   148  
   149  func (p *DynamicTaskPool[Key]) ExecuteSuccess() int {
   150  	return int(atomic.LoadUint64(&p.execSuccess))
   151  }
   152  
   153  func (p *DynamicTaskPool[Key]) ExecuteError() int {
   154  	return int(atomic.LoadUint64(&p.execFailed))
   155  }