github.com/bytedance/gopkg@v0.0.0-20240514070511-01b2cbcf35e1/util/gopool/pool.go (about)

     1  // Copyright 2021 ByteDance Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package gopool
    16  
    17  import (
    18  	"context"
    19  	"sync"
    20  	"sync/atomic"
    21  )
    22  
    23  type Pool interface {
    24  	// Name returns the corresponding pool name.
    25  	Name() string
    26  	// SetCap sets the goroutine capacity of the pool.
    27  	SetCap(cap int32)
    28  	// Go executes f.
    29  	Go(f func())
    30  	// CtxGo executes f and accepts the context.
    31  	CtxGo(ctx context.Context, f func())
    32  	// SetPanicHandler sets the panic handler.
    33  	SetPanicHandler(f func(context.Context, interface{}))
    34  	// WorkerCount returns the number of running workers
    35  	WorkerCount() int32
    36  }
    37  
    38  var taskPool sync.Pool
    39  
    40  func init() {
    41  	taskPool.New = newTask
    42  }
    43  
    44  type task struct {
    45  	ctx context.Context
    46  	f   func()
    47  
    48  	next *task
    49  }
    50  
    51  func (t *task) zero() {
    52  	t.ctx = nil
    53  	t.f = nil
    54  	t.next = nil
    55  }
    56  
    57  func (t *task) Recycle() {
    58  	t.zero()
    59  	taskPool.Put(t)
    60  }
    61  
    62  func newTask() interface{} {
    63  	return &task{}
    64  }
    65  
    66  type taskList struct {
    67  	sync.Mutex
    68  	taskHead *task
    69  	taskTail *task
    70  }
    71  
    72  type pool struct {
    73  	// The name of the pool
    74  	name string
    75  
    76  	// capacity of the pool, the maximum number of goroutines that are actually working
    77  	cap int32
    78  	// Configuration information
    79  	config *Config
    80  	// linked list of tasks
    81  	taskHead  *task
    82  	taskTail  *task
    83  	taskLock  sync.Mutex
    84  	taskCount int32
    85  
    86  	// Record the number of running workers
    87  	workerCount int32
    88  
    89  	// This method will be called when the worker panic
    90  	panicHandler func(context.Context, interface{})
    91  }
    92  
    93  // NewPool creates a new pool with the given name, cap and config.
    94  func NewPool(name string, cap int32, config *Config) Pool {
    95  	p := &pool{
    96  		name:   name,
    97  		cap:    cap,
    98  		config: config,
    99  	}
   100  	return p
   101  }
   102  
   103  func (p *pool) Name() string {
   104  	return p.name
   105  }
   106  
   107  func (p *pool) SetCap(cap int32) {
   108  	atomic.StoreInt32(&p.cap, cap)
   109  }
   110  
   111  func (p *pool) Go(f func()) {
   112  	p.CtxGo(context.Background(), f)
   113  }
   114  
   115  func (p *pool) CtxGo(ctx context.Context, f func()) {
   116  	t := taskPool.Get().(*task)
   117  	t.ctx = ctx
   118  	t.f = f
   119  	p.taskLock.Lock()
   120  	if p.taskHead == nil {
   121  		p.taskHead = t
   122  		p.taskTail = t
   123  	} else {
   124  		p.taskTail.next = t
   125  		p.taskTail = t
   126  	}
   127  	p.taskLock.Unlock()
   128  	atomic.AddInt32(&p.taskCount, 1)
   129  	// The following two conditions are met:
   130  	// 1. the number of tasks is greater than the threshold.
   131  	// 2. The current number of workers is less than the upper limit p.cap.
   132  	// or there are currently no workers.
   133  	if (atomic.LoadInt32(&p.taskCount) >= p.config.ScaleThreshold && p.WorkerCount() < atomic.LoadInt32(&p.cap)) || p.WorkerCount() == 0 {
   134  		p.incWorkerCount()
   135  		w := workerPool.Get().(*worker)
   136  		w.pool = p
   137  		w.run()
   138  	}
   139  }
   140  
   141  // SetPanicHandler the func here will be called after the panic has been recovered.
   142  func (p *pool) SetPanicHandler(f func(context.Context, interface{})) {
   143  	p.panicHandler = f
   144  }
   145  
   146  func (p *pool) WorkerCount() int32 {
   147  	return atomic.LoadInt32(&p.workerCount)
   148  }
   149  
   150  func (p *pool) incWorkerCount() {
   151  	atomic.AddInt32(&p.workerCount, 1)
   152  }
   153  
   154  func (p *pool) decWorkerCount() {
   155  	atomic.AddInt32(&p.workerCount, -1)
   156  }